From 53912b2f99aeaaac7cf7d2a2cefca9ec8bc51de1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 13:31:31 +0000 Subject: [PATCH] Studio: tighten image-generation fix comments and docstrings --- scripts/uninstall.ps1 | 8 +- scripts/uninstall.sh | 15 ++-- studio/backend/core/inference/diffusion.py | 11 +-- .../core/inference/diffusion_families.py | 12 +-- .../backend/core/inference/diffusion_lora.py | 14 ++- studio/backend/core/inference/gpu_arbiter.py | 23 ++--- .../backend/core/inference/image_gallery.py | 10 +-- .../backend/core/inference/sd_cpp_backend.py | 27 +++--- .../backend/core/inference/video_gallery.py | 25 +++--- .../core/training/diffusion_train_common.py | 5 +- studio/backend/routes/inference.py | 28 +++--- studio/backend/routes/models.py | 89 ++++++++----------- studio/backend/routes/training.py | 48 ++++------ studio/backend/routes/video.py | 17 ++-- .../backend/tests/test_cached_gguf_routes.py | 11 +-- studio/backend/tests/test_gpu_arbiter.py | 10 +-- .../backend/tests/test_local_model_format.py | 27 +++--- studio/backend/tests/test_sd_cpp_backend.py | 15 ++-- studio/backend/tests/test_video_gallery.py | 5 +- .../src/features/images/images-page.tsx | 7 +- 20 files changed, 161 insertions(+), 246 deletions(-) diff --git a/scripts/uninstall.ps1 b/scripts/uninstall.ps1 index 46eb48cb3b..453796bf60 100644 --- a/scripts/uninstall.ps1 +++ b/scripts/uninstall.ps1 @@ -370,11 +370,9 @@ function Uninstall-UnslothStudio { if ($defaultSdCpp -and (Test-Path -LiteralPath $defaultSdCpp) -and (Test-Path -LiteralPath (Join-Path $defaultSdCpp ".unsloth-studio-owned") -PathType Leaf)) { $defaultSdCppToStop = $defaultSdCpp } - # Custom/env-mode sd.cpp builds sit BESIDE each custom root at \stable-diffusion.cpp - # (find_sd_cpp_binary resolves from UNSLOTH_STUDIO_HOME.parent), so a running owned sd-server - # there is outside every root above ($knownRoots holds the custom root, not its sibling). We - # delete those marker-owned dirs below, so add them to the handle scan too -- gated on the same - # owner marker as the delete, matching the default-root handling. + # Custom/env-mode sd.cpp builds sit BESIDE each custom root at \stable-diffusion.cpp, + # outside $knownRoots. We delete those marker-owned dirs below, so add them to the handle scan + # too, gated on the same owner marker. $customSdCppToStop = @() foreach ($r in $customRoots) { $sdc = Join-Path (Split-Path -LiteralPath $r -Parent) "stable-diffusion.cpp" diff --git a/scripts/uninstall.sh b/scripts/uninstall.sh index 5f1327b994..3e0b40842b 100755 --- a/scripts/uninstall.sh +++ b/scripts/uninstall.sh @@ -35,12 +35,9 @@ _pkill_escape() { printf '%s' "$1" | sed -e 's:[][\\.^$*+?{|}()/]:\\&:g' } -# Owned sd.cpp roots (default + custom siblings), each gated on the install-time -# owner marker. Native diffusion builds beside a custom/env root at -# /stable-diffusion.cpp (find_sd_cpp_binary resolves from -# UNSLOTH_STUDIO_HOME.parent) and at $HOME/.unsloth/stable-diffusion.cpp by default. -# The marker is mandatory so we never stop a user-managed sd-server from an -# unrelated checkout that happens to sit at one of these paths. +# Owned sd.cpp roots (default $HOME/.unsloth/stable-diffusion.cpp + each custom root's +# /stable-diffusion.cpp sibling), each gated on the install-time owner marker so we never +# stop a user-managed sd-server from an unrelated checkout at one of these paths. _owned_sd_cpp_roots() { _default_sd="$HOME/.unsloth/stable-diffusion.cpp" [ -f "$_default_sd/.unsloth-studio-owned" ] && printf '%s\n' "$_default_sd" @@ -51,10 +48,8 @@ _owned_sd_cpp_roots() { done } -# pkill resident sd-server / sd-cli whose executable lives under an owned sd.cpp -# root, BEFORE that tree is removed below: a live native server keeps running -# after its binary is unlinked. Anchored on the owned root so an unrelated -# checkout's sd-server is never matched. +# pkill resident sd-server / sd-cli under an owned sd.cpp root before that tree is removed (a live +# native server keeps running after its binary is unlinked). Anchored on the owned root. _stop_owned_sd_cpp_processes() { _signal="$1" command -v pkill >/dev/null 2>&1 || return 0 diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 6ec6b19c2b..073feb6305 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1908,9 +1908,8 @@ class DiffusionBackend: # resolve_controlnet accepts a bare owner/name without the base trust gate, and # from_pretrained deserializes it (a malicious pickle would execute), so run the same # Hub malware preflight the chat/export loaders use. A local dir is exempt (fail-open). - # The preflight fails OPEN when the Hub scan is unavailable (offline / missing metadata), - # so for a remote repo also force safetensors below: that closes the pickle RCE vector - # even when the scan could not run. + # The preflight fails OPEN when the Hub scan is unavailable, so a remote repo also forces + # safetensors below, closing the pickle RCE vector even when the scan could not run. remote_cn = not getattr(resolved_cn, "is_local", False) if remote_cn: from utils.security import evaluate_file_security @@ -1928,10 +1927,8 @@ class DiffusionBackend: # state.dtype is the display string ("bfloat16"), not a torch.dtype; pass the real # dtype so diffusers loads at the base compute dtype, not float32 (extra VRAM). cn_dtype = getattr(torch, str(state.dtype).replace("torch.", ""), None) - # Force safetensors for an untrusted remote repo: a bare owner/name reaches here without - # the base trust gate, and if the Hub scan failed open above, an embedded pickle would - # still deserialize on load. Requiring safetensors refuses that vector (curated - # ControlNets are all safetensors). A local dir the user chose is exempt. + # Force safetensors for an untrusted remote repo: if the Hub scan failed open above, an + # embedded pickle would still deserialize on load. A local dir the user chose is exempt. cn_from_pretrained_kwargs: dict[str, Any] = {} if remote_cn: cn_from_pretrained_kwargs["use_safetensors"] = True diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index d116999d39..f345d36fc0 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -470,10 +470,8 @@ def family_sd_cpp_supported(fam: DiffusionFamily) -> bool: return bool(fam.sd_cpp_vae and fam.sd_cpp_text_encoders) -# FLUX.2-klein ships two variants that use DIFFERENT text encoders: the 4B transformer pairs with -# Qwen3-4B and the 9B transformer pairs with Qwen3-8B (Black Forest Labs model cards; a mismatched -# encoder fails with tensor-shape / dtype errors deep in sd-cli). The family table carries only one -# default, so the native (sd.cpp) encoder is selected per variant from the load identity below. +# FLUX.2-klein's 9B transformer pairs with Qwen3-8B, the 4B with the family-default Qwen3-4B +# (a mismatched encoder fails deep in sd-cli), so the encoder is picked per variant below. _FLUX2_KLEIN_9B_SD_CPP_TEXT_ENCODERS = ( ( "Comfy-Org/vae-text-encorder-for-flux-klein-9b", @@ -490,10 +488,8 @@ def sd_cpp_text_encoders_for( ) -> tuple[tuple[str, str, str], ...]: """The sd.cpp text encoders for a specific load. - FLUX.2-klein selects its encoder by variant (the 9B transformer needs Qwen3-8B, the 4B needs the - family default Qwen3-4B); every other family returns its static table. Keyed on the load identity - (repo id + GGUF filename) so ``unsloth/FLUX.2-klein-9B-GGUF`` and a local ``*klein-9B*.gguf`` both - resolve to the 8B encoder.""" + FLUX.2-klein picks by variant (9B needs Qwen3-8B, 4B the family default) keyed on the load + identity (repo id + GGUF filename); every other family returns its static table.""" if fam.name == "flux.2-klein": identity = f"{repo_id or ''}/{gguf_filename or ''}".lower() if "klein-9b" in identity or "klein_9b" in identity: diff --git a/studio/backend/core/inference/diffusion_lora.py b/studio/backend/core/inference/diffusion_lora.py index 85caa4f719..ecaeb688a7 100644 --- a/studio/backend/core/inference/diffusion_lora.py +++ b/studio/backend/core/inference/diffusion_lora.py @@ -213,10 +213,9 @@ def resolve_one( Downloads hub weights via the xet-fallback helper. Raises FileNotFoundError/ValueError on an unresolvable/unsupported id, which the caller maps to a 400. - ``family`` (the loaded model family) enforces catalog family tags HERE, not only in the picker - (``list_loras``): a LoRA is architecture-specific, so a direct API client sending an id tagged - for another family would otherwise load it through the wrong pipeline. Mirrors the ControlNet - resolver's family gate. An untagged catalog entry (empty ``families``) stays unrestricted. + ``family`` (the loaded model family) enforces catalog family tags HERE, not only in the picker, + so a direct API client cannot load a LoRA tagged for another family through the wrong pipeline. + An untagged catalog entry (empty ``families``) stays unrestricted. """ # An empty/whitespace token triggers an auth error instead of anonymous access; normalise to None. hf_token = hf_token.strip() if hf_token and hf_token.strip() else None @@ -307,10 +306,9 @@ def resolve_specs( ) -> list[ResolvedLora]: """Resolve request (id, weight) pairs, dropping zero-weight entries. - ``family`` is the loaded model family; it enforces catalog family tags in :func:`resolve_one` - for direct API callers, not only the UI picker. Maps the named not-found/gated Hub errors to a - 400 (URL scrubbed); does NOT catch the base HfHubHTTPError, so a Hub 5xx stays a 500. A - mid-download cancel maps to a 409.""" + ``family`` (the loaded model family) enforces catalog family tags in :func:`resolve_one` for + direct API callers. Maps the named not-found/gated Hub errors to a 400 (URL scrubbed); does NOT + catch the base HfHubHTTPError, so a Hub 5xx stays a 500. A mid-download cancel maps to a 409.""" from huggingface_hub.errors import ( EntryNotFoundError, GatedRepoError, diff --git a/studio/backend/core/inference/gpu_arbiter.py b/studio/backend/core/inference/gpu_arbiter.py index 0a53701a0f..04a7cd08a0 100644 --- a/studio/backend/core/inference/gpu_arbiter.py +++ b/studio/backend/core/inference/gpu_arbiter.py @@ -67,13 +67,11 @@ _EVICTORS = {CHAT: _evict_chat, DIFFUSION: _evict_diffusion, VIDEO: _evict_video def acquire_for(owner: str, register: Optional[Callable[[], Any]] = None) -> Any: """Make ``owner`` the sole GPU owner, evicting the other if it holds it. - ``register``, if given, runs under the arbiter lock right after ownership transfers, - and its return value is returned. Registering the in-flight load HERE -- not after - ``acquire_for`` returns -- closes the window where a competing acquire could evict this - owner before its load is marked in-flight: eviction would then find nothing to cancel - and both loaders would allocate VRAM at once. ``register`` must be quick (it holds the - lock) and must not re-enter the arbiter. If it raises, ownership stays with ``owner`` -- - matching the pre-register behaviour where a failed load left the handoff in place. + ``register``, if given, runs under the arbiter lock right after ownership transfers and its + return value is returned. Marking the in-flight load HERE (not after ``acquire_for`` returns) + closes the window where a competing acquire could evict this owner before its load is in-flight, + letting both loaders allocate VRAM at once. It must be quick and not re-enter the arbiter; if it + raises, ownership stays with ``owner``. """ global _owner if owner not in _EVICTORS: @@ -97,13 +95,10 @@ def release(owner: str) -> None: def release_if(owner: str, predicate: Callable[[], bool]) -> bool: """Drop ``owner``'s claim only if it still holds it AND ``predicate()`` is true, atomically. - A slow unload's "nothing resident / no load in flight" check and the ``release`` must not - straddle a concurrent same-owner load: that load's ``acquire_for(register=...)`` re-registers - ownership UNDER this lock, so a plain check-then-``release`` could pass the stale check and then - clear the newer claim (``release`` is owner-guarded but identity-less). Evaluating the predicate - under the lock closes that window -- the load's register runs either fully before or fully after. - ``predicate`` must be quick (it holds the lock) and must not re-enter the arbiter. Returns True - iff ownership was dropped.""" + A slow unload's idle check and its ``release`` must not straddle a concurrent same-owner load + whose ``acquire_for(register=...)`` re-registers ownership under this lock; evaluating the + predicate under the lock keeps them atomic so ``release`` never clears the newer claim. + ``predicate`` must be quick and not re-enter the arbiter. Returns True iff ownership was dropped.""" global _owner with _lock: if _owner != owner or not predicate(): diff --git a/studio/backend/core/inference/image_gallery.py b/studio/backend/core/inference/image_gallery.py index fee2eb1b72..77c4b88e58 100644 --- a/studio/backend/core/inference/image_gallery.py +++ b/studio/backend/core/inference/image_gallery.py @@ -188,9 +188,8 @@ def delete(image_id: str) -> bool: path = image_path(image_id) if path is None: return False - # Only delete files we actually own (a readable recipe chunk). A foreign PNG that a caller - # dropped in by hand is invisible to list_images, so deleting it here on a guessed id would - # silently destroy a file the gallery never claimed. + # Only delete files we own (a readable recipe chunk); a hand-dropped foreign PNG is invisible + # to list_images, so a guessed id must not destroy it. if _read_meta(path) is None: return False try: @@ -202,10 +201,9 @@ def delete(image_id: str) -> bool: def clear() -> int: - """Delete every Studio-owned gallery PNG; return how many were removed. + """Delete every Studio-owned gallery PNG (readable recipe chunk); return how many were removed. - Preserves foreign PNGs (no readable recipe chunk): list_images already hides them, so clear - must not silently destroy files the gallery never surfaced.""" + Foreign PNGs are preserved: list_images already hides them, so clear must not destroy them.""" removed = 0 try: paths = list(gallery_dir().glob("*.png")) diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index b22ebee396..5d4a5aa4c0 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -220,10 +220,9 @@ class _SdState: mode: str = "server" # Token kept so LoRA adapters selected at generate time can be fetched from the Hub. hf_token: Optional[str] = None - # The single-file GGUF basename this load committed. Kept so companion resolution - # (sd_cpp_text_encoders_for) reproduces the load identity -- some variants pick their - # encoder by filename (FLUX.2-klein-9B -> Qwen3-8B), and a local *klein-9B*.gguf carries - # that keyword only in the basename, not the repo id. + # The GGUF basename this load committed, so companion resolution reproduces the load identity: + # some variants pick their encoder by filename (FLUX.2-klein-9B -> Qwen3-8B) and a local + # *klein-9B*.gguf carries that keyword only in the basename, not the repo id. gguf_filename: Optional[str] = None @@ -606,8 +605,7 @@ class SdCppDiffusionBackend: specs: list[tuple[str, str, str]] = [(repo_id, gguf_filename, "diffusion_model")] if fam.sd_cpp_vae: specs.append((fam.sd_cpp_vae[0], fam.sd_cpp_vae[1], "vae")) - # Select the text encoder per variant (FLUX.2-klein 4B->Qwen3-4B, 9B->Qwen3-8B) from the - # load identity, not the family's single default, so a 9B GGUF fetches the right encoder. + # Pick the encoder per variant from the load identity so a 9B GGUF fetches the right one. for terepo, tefile, kind in sd_cpp_text_encoders_for(fam, repo_id, gguf_filename): specs.append((terepo, tefile, kind)) return specs @@ -703,11 +701,9 @@ class SdCppDiffusionBackend: repos = [state.repo_id, state.base_repo] if fam.sd_cpp_vae: repos.append(fam.sd_cpp_vae[0]) - # Same per-variant encoder selection as _asset_specs, keyed on the loaded repo id AND - # GGUF filename, so the cache-deletion guard protects the encoder repo this load actually - # downloaded (the 9B variant's Qwen3-8B, not the 4B default) -- a local *klein-9B*.gguf - # carries that keyword only in the basename, so dropping the filename would fall back to - # the 4B default and protect the wrong repo. + # Same per-variant selection as _asset_specs (keyed on repo id AND GGUF filename) so the + # cache-deletion guard protects the encoder repo this load actually downloaded; dropping + # the filename would fall back to the 4B default and protect the wrong repo. repos.extend( terepo for terepo, _f, _k in sd_cpp_text_encoders_for( @@ -783,11 +779,10 @@ class SdCppDiffusionBackend: self._state = None raise RuntimeError(DIFFUSION_NOT_LOADED_MSG) self._active_generate_cancel = cancel - # Publish an active (step 0) state now, before the slow pre-generate setup - # (LoRA listing/download), so a reload's progress probe doesn't read idle - # while this generation already holds _generate_lock and let a second generate - # queue behind it. The parsed sd-cli progress lines advance this step count. - # Mirrors DiffusionBackend.generate, which publishes _gen before its setup. + # Publish an active (step 0) state before the slow pre-generate setup (LoRA + # listing/download) so a reload's progress probe doesn't read idle while this + # generation holds _generate_lock and let a second generate queue behind it. + # Mirrors DiffusionBackend.generate; sd-cli progress lines advance this count. self._gen = _SdGen(total_steps = int(steps)) try: if seed is None: diff --git a/studio/backend/core/inference/video_gallery.py b/studio/backend/core/inference/video_gallery.py index 292d0a72b8..1cddb94b7b 100644 --- a/studio/backend/core/inference/video_gallery.py +++ b/studio/backend/core/inference/video_gallery.py @@ -172,12 +172,10 @@ def _sidecar_path(video_id: str) -> Path: return gallery_dir() / f"{video_id}.json" -# The sidecar keys a genuine Studio record always carries (save() always writes them). delete() and -# clear() treat a pair as owned only when its sidecar has all of these, so a hand-dropped MP4 with a -# parseable-but-empty ("{}") or partial JSON sidecar -- which list_videos already hides via the -# GalleryVideo schema filter -- is neither counted as ours nor destroyed. Mirrors -# image_gallery._REQUIRED_META: a key-presence check (the route still owns full schema/value-type -# validation), aligned with GalleryVideo's required stored fields. +# Sidecar keys every genuine Studio record carries (save() always writes them). delete()/clear() +# own a pair only when its sidecar has all of these, so a hand-dropped MP4 with an empty ("{}") or +# partial sidecar -- which list_videos already hides -- is neither counted as ours nor destroyed. +# Key-presence only (the route owns full schema validation); mirrors image_gallery._REQUIRED_META. _REQUIRED_META = ( "prompt", "width", @@ -201,9 +199,8 @@ def _read_meta(sidecar: Path) -> Optional[dict[str, Any]]: meta = json.loads(raw) except (ValueError, TypeError): return None - # A parseable dict is not enough to claim ownership: a foreign sidecar (e.g. "{}") or one from a - # different schema lacks these keys. Require them so delete()/clear() never destroy a clip the - # gallery never surfaced (mirrors image_gallery._read_meta). + # A parseable dict is not enough: a foreign ("{}") or different-schema sidecar lacks these keys. + # Require them so delete()/clear() never destroy a clip the gallery never surfaced. if not isinstance(meta, dict) or any(k not in meta for k in _REQUIRED_META): return None return meta @@ -258,9 +255,8 @@ def delete(video_id: str) -> bool: path = video_path(video_id) if path is None: return False - # Only delete a pair we actually own (a readable sidecar). A foreign / orphan MP4 is invisible - # to list_videos, so deleting it here on a guessed id would silently destroy a file the gallery - # never claimed. + # Only delete a pair we own (a readable sidecar); a foreign/orphan MP4 is invisible to + # list_videos, so a guessed id must not destroy it. if _read_meta(_sidecar_path(video_id)) is None: return False # Delete the MP4 FIRST: if the sidecar were dropped first and the mp4 unlink then failed (lock / @@ -280,10 +276,9 @@ def delete(video_id: str) -> bool: def clear() -> int: - """Delete every Studio-owned gallery pair; return how many videos were removed. + """Delete every Studio-owned gallery pair (readable sidecar); return how many were removed. - Preserves foreign / orphan MP4s (no readable sidecar): list_videos already hides them, so clear - must not silently destroy files the gallery never surfaced.""" + Foreign/orphan MP4s are preserved: list_videos already hides them, so clear must not destroy them.""" removed = 0 try: paths = list(gallery_dir().glob("*.mp4")) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index ae8c3c849e..959709c2a1 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -688,9 +688,8 @@ def discover_image_caption_pairs( meta_path = root / meta_name if not meta_path.is_file(): continue - # Tolerate a bad upload: invalid UTF-8 in the file, or a line that is valid JSON but not an - # object (``[]`` / ``null`` / a string / a number). Neither should crash the trainer -- the - # record is simply skipped so the instance_prompt fallback still applies. + # Tolerate a bad upload (invalid UTF-8, or a line of non-object JSON): skip the record so + # the instance_prompt fallback still applies rather than crashing the trainer. try: meta_lines = meta_path.read_text(encoding = "utf-8").splitlines() except (OSError, UnicodeError): diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 50d6a5f5fe..9807f2af60 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -14340,8 +14340,8 @@ async def load_diffusion_model( needs_gpu = device != "cpu" def _begin_load(): - # Kicks the (slow) load onto a background thread and returns at once (the client - # polls images/load-progress); begin_load itself validates network-free. + # Kicks the (slow) load onto a background thread and returns at once (client polls + # images/load-progress); begin_load itself validates network-free. return engine.begin_load( request.model_path, gguf_filename = request.gguf_filename, @@ -14362,11 +14362,9 @@ async def load_diffusion_model( ) if needs_gpu: - # Register the in-flight load UNDER the arbiter lock (not after acquire_for - # returns): a competing Video/chat acquire in that gap would otherwise evict - # DIFFUSION before begin_load marks a load in-flight, so eviction finds nothing - # to cancel and both loaders allocate VRAM at once. begin_load returns at once, - # so the lock is held only briefly. + # Register the in-flight load UNDER the arbiter lock (not after acquire_for returns): + # otherwise a competing Video/chat acquire in that gap evicts DIFFUSION before the load + # is marked in-flight, finds nothing to cancel, and both loaders allocate VRAM at once. status_dict = await asyncio.to_thread(acquire_for, DIFFUSION, _begin_load) else: # A CPU-only native load never touches the GPU, so it neither acquires nor is @@ -14474,11 +14472,9 @@ async def generate_diffusion_image( "steps": request.steps, "guidance": request.guidance, "seed": seed, - # The base seed the batch launched with. The native engine derives per-image + # Base seed the batch launched with. The native engine derives per-image # seeds as base + index, so ``seed`` above is already advanced for index>0; - # restore must replay from this base (with batch_size) or it would advance a - # second time and reproduce a different image. Diffusers shares one seed, so - # base == seed there. + # restore replays from this base (diffusers shares one seed, so base == seed). "batch_seed": result["seed"], # Position within the batch (shared timestamp), so the export filename # stays unique. @@ -14589,12 +14585,10 @@ async def unload_diffusion_model(current_subject: str = Depends(get_current_subj status_dict = await asyncio.to_thread(get_active_diffusion_engine().unload) # Drop DIFFUSION ownership only if nothing is resident AND no new load is in flight: a # concurrent /images/load that re-acquired DIFFUSION while this (slow) unload ran must keep - # ownership, or a later chat load would see no owner, skip eviction, and OOM the newly - # resident pipeline. An in-flight load has is_loaded False for its whole download/finalize - # window, so gate on loading_repo_ids() too, not just committed state. The idle check and the - # release must be ATOMIC (release_if): the load's acquire_for register runs under the same - # arbiter lock, so a plain check-then-release could pass the stale check and then clear the - # newer claim. + # ownership, or a later chat load sees no owner, skips eviction, and OOMs the newly resident + # pipeline. An in-flight load has is_loaded False for its whole window, so gate on + # loading_repo_ids() too. The idle check and release must be ATOMIC (release_if): the load's + # register runs under the same lock, so a plain check-then-release could clear the newer claim. engine = get_active_diffusion_engine() await asyncio.to_thread( release_if, diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index d8226473ce..cf728b46d7 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -305,9 +305,8 @@ def _has_non_gguf_weights(path: Path) -> bool: def _local_pipeline_index(d: Path) -> bool: - """True when *d* is a standard diffusers PIPELINE root: component weights/configs live in - subdirs (``transformer/``, ``vae/``, ...) under a top-level ``model_index.json``, so - ``_is_model_directory`` (which wants a root config + loose weights) rejects it.""" + """True when *d* is a diffusers PIPELINE root (top-level ``model_index.json``, weights in + component subdirs), which ``_is_model_directory`` (root config + loose weights) rejects.""" try: return (d / "model_index.json").is_file() except OSError: @@ -318,11 +317,10 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca if not models_dir.exists() or not models_dir.is_dir(): return [] - # A scan folder can point directly at a diffusers PIPELINE dir, not only at a parent of - # model repos. _is_model_directory rejects such a root (weights live in transformer/, vae/, - # ... not beside a root config.json), so without this the child scan below surfaces the - # component subdirs as bogus models and hides the real pipeline. The Images/Video load path - # loads a local pipeline dir, so admit the root as one model (task tagging classifies it). + # A scan folder can point directly at a diffusers PIPELINE dir, which _is_model_directory + # rejects; without admitting it the child scan surfaces the component subdirs as bogus models + # and hides the real pipeline. The Images/Video load path loads it, so admit the root as one + # model (task tagging classifies it). _is_self_model = _is_model_directory(models_dir) or _local_pipeline_index(models_dir) if _is_self_model: @@ -353,11 +351,9 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca has_config = (child / "config.json").exists() or ( child / "adapter_config.json" ).exists() - # A standard diffusers PIPELINE folder keeps its weights/configs in component - # subdirs (transformer/, vae/, ...) and carries only model_index.json at the - # root, so the checks above miss it. The Images/Video load path accepts such a - # local pipeline dir, so admit it here too (task tagging then classifies it via - # _local_is_diffusers); otherwise it is hidden from the On Device picker. + # A diffusers PIPELINE folder (weights in component subdirs, only model_index.json at + # the root) is missed by the checks above; the Images/Video load path accepts it, so + # admit it too or it is hidden from the On Device picker. has_pipeline_index = _local_pipeline_index(child) has_model_files = has_gguf or has_non_gguf_weights or has_config or has_pipeline_index except OSError: @@ -3277,13 +3273,11 @@ def _repo_gguf_task(repo_info) -> Optional[str]: def _local_family_needles(model: "LocalModelInfo") -> tuple[str, ...]: - """Family-detection hints for a local (non-GGUF) checkpoint: its model id, display name, and - leaf directory name, plus -- for a bare single-file directory -- the sole checkpoint's - filename. A generically named folder holding one loadable ``qwen-image-*.safetensors`` / - ``ltx-*.safetensors`` identifies its family only from that filename, and the load route already - resolves that sole file via ``resolve_local_single_file``, so feed the same name here or a - task-scoped Images/Video picker (which rejects ``task: null``) hides the on-device model. Only - the basename is used (not the parent path), so a family token in a parent dir can't match.""" + """Family-detection hints for a local (non-GGUF) checkpoint: model id, display name, leaf dir + name, and -- for a bare single-file dir -- the sole checkpoint's filename (a generic folder + holding one ``qwen-image-*.safetensors`` identifies its family only there, and the load route + resolves it via ``resolve_local_single_file``). Only basenames, so a parent-dir token can't + match.""" needles = [model.model_id, model.display_name, Path(model.id).name] try: from core.inference.diffusion import resolve_local_single_file @@ -3332,14 +3326,10 @@ def _local_model_task(model: "LocalModelInfo") -> Optional[str]: return _VIDEO_GEN_TASK except Exception: pass - # The Images load path resolves the family via detect_family_for_pick and REJECTS a pick - # whose id / name / checkpoint filename carries no supported image-family token - # (diffusion.py validate_load_request), 400ing AFTER it has already evicted the GPU owner. - # A bare model_index.json directory alone (a generically named on-device pipeline) is not - # enough. Tag text-to-image only when that same family detection succeeds, so the picker - # never advertises a local pipeline the load will always reject. Detection uses the same - # _local_family_needles the video branch does (leaf name / id / sole-file, not the raw - # path), so a family token in a parent directory can't spuriously tag it. + # The Images load path rejects a pick with no supported image-family token, 400ing AFTER + # evicting the GPU owner (a bare model_index.json dir is not enough). Tag text-to-image + # only when that same detection succeeds, so the picker never advertises a pipeline the + # load will always reject. try: from core.inference.diffusion_families import detect_family for needle in _local_family_needles(model): @@ -3347,8 +3337,8 @@ def _local_model_task(model: "LocalModelInfo") -> Optional[str]: return "text-to-image" return None except Exception: - # Detection unavailable (import/exec error): fall back to the prior permissive tag - # rather than hiding a possibly-loadable pipeline. + # Detection unavailable: fall back to the prior permissive tag rather than hiding a + # possibly-loadable pipeline. return "text-to-image" return None @@ -3358,9 +3348,8 @@ def _local_is_diffusers(model: "LocalModelInfo") -> bool: ``_repo_is_diffusers`` heuristics: a full pipeline carries a top-level ``model_index.json``, while single-file / safetensors image checkpoints ship none, so fall back to the model id resolving to a known diffusion family (the same resolver the - Images backend loads from). Family detection uses the clean model id / name and the sole - checkpoint's filename (via _local_family_needles), not the on-disk path, so a parent - directory keyword can't spuriously match while a filename-only family is still caught.""" + Images backend loads from). Family detection uses _local_family_needles (id / name / sole + checkpoint filename, not the on-disk path), so a parent-dir keyword can't spuriously match.""" try: p = Path(model.path) if p.is_dir() and (p / "model_index.json").is_file(): @@ -3374,12 +3363,9 @@ 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) has no - # pipeline index and no image family, so the checks above miss it. The video load route loads it - # as a single_file (routes/video.py), so it must be surfaced or _local_model_task returns - # task=null and the picker hides it. Match clean id / name / checkpoint-filename needles (not - # the raw path) so a parent-dir token can't spuriously match; _local_model_task then routes it - # to text-to-video. + # A single-file VIDEO checkpoint (LTX / Wan / Hunyuan .safetensors, no model_index.json) is + # missed above but loaded as a single_file by the video route, so surface it or the picker + # hides it. Uses _local_family_needles; _local_model_task then routes it to text-to-video. try: from core.inference.video_families import detect_video_family for needle in _local_family_needles(model): @@ -3488,15 +3474,12 @@ def _repo_is_diffusers(repo_info) -> bool: def _repo_pipeline_missing_denoiser(repo_info) -> bool: - """True for a diffusers-pipeline snapshot (root ``model_index.json``) whose denoiser - component (``transformer/`` or ``unet/``) carries NO weight file. This is the shape of a - companion-only prefetch: a GGUF image load pulls the base repo's VAE / text-encoder / - ``model_index.json`` into the cache but deliberately skips the multi-GB transformer (the GGUF - supplies it), so the snapshot has a pipeline manifest yet is not a loadable BF16 pipeline -- - ``from_pretrained`` on it re-downloads the missing shards. ``_cached_repo_partial`` misses this - (no cancel marker / .incomplete blob, and hf_hub_download writes no manifest), so ``/cached-models`` - would advertise it as fully on-device. The caller marks such rows partial. Best-effort: any scan - error reports not-missing so a glitch never hides a genuinely complete pipeline.""" + """True for a diffusers-pipeline snapshot (root ``model_index.json``) whose denoiser component + (``transformer/`` or ``unet/``) carries NO weight file -- the shape of a companion-only prefetch + where a GGUF image load pulled the base repo's VAE / text-encoder / manifest but skipped the + multi-GB transformer (the GGUF supplies it). ``_cached_repo_partial`` misses this, so the caller + marks such rows partial. Best-effort: any scan error reports not-missing so a glitch never hides + a genuinely complete pipeline.""" if not _repo_has_pipeline_index(repo_info): return False _DENOISER_DIRS = ("transformer", "unet") @@ -3514,8 +3497,8 @@ def _repo_pipeline_missing_denoiser(repo_info) -> bool: except ValueError: parts = () if not parts: - # No snapshot scoping (or file outside it): fall back to the recorded name, - # which may itself carry the component subdir (e.g. 'transformer/model...'). + # No snapshot scoping: fall back to the recorded name, which may itself carry + # the component subdir (e.g. 'transformer/model...'). parts = Path(name).parts if ( len(parts) >= 2 @@ -3607,9 +3590,9 @@ async def list_cached_models( ) key = repo_id.lower() existing = seen_lower.get(key) - # A companion-only prefetch (root model_index.json + VAE / text-encoder but no - # transformer/ shards, pulled to back a GGUF load) is not a loadable BF16 - # pipeline; treat it as partial so the picker does not advertise it as on-device. + # A companion-only prefetch (manifest + VAE/text-encoder but no transformer + # shards) is not a loadable pipeline; treat it as partial so the picker does + # not advertise it as on-device. is_partial = _cached_repo_partial( repo_id, Path(repo_info.repo_path) ) or _repo_pipeline_missing_denoiser(repo_info) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 04d62f2dc9..1173fb21b0 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1140,10 +1140,9 @@ def _diffusion_training_active() -> bool: def _require_diffusion_dataset_mutable() -> None: """Reject a dataset mutation while a diffusion run is active. - The trainer enumerates and (when the latent cache is off / over budget) re-opens dataset images - during the loop, so uploading, importing, captioning, or deleting underneath it makes the run - nondeterministic or raises a FileNotFoundError mid-step. Best-effort: a service-import failure - fails open (never blocks a mutation on an unknowable state), matching the start interlock.""" + The trainer re-opens dataset images during the loop, so mutating underneath it makes the run + nondeterministic or raises a FileNotFoundError mid-step. Fails open (a service-import failure + never blocks a mutation on an unknowable state), matching the start interlock.""" if _diffusion_training_active(): raise HTTPException( status_code = 409, @@ -1271,11 +1270,9 @@ def _resolve_diffusion_data_dir(raw: str) -> Path: # Single component and not ".." -> joining under datasets_root() cannot escape it. if not p.is_absolute() and len(p.parts) == 1 and p.parts[0] != "..": direct = datasets_root() / value - # Route a bare image-dataset name through the SAME protected resolver the - # caption/delete/labeling CRUD routes use, so a name -> external-directory symlink is - # rejected here too (is_dir() follows the link, so a plain is_dir() check would train on - # files outside the datasets root). Include a broken symlink so it is rejected, not - # silently passed through to resolve_dataset_path. + # Route a bare name through the same protected resolver the CRUD routes use, so a + # name -> external-directory symlink is rejected here too (is_dir() follows the link). + # Include a broken symlink so it is rejected, not passed to resolve_dataset_path. if direct.is_dir() or direct.is_symlink(): return _resolve_dataset_folder(value) return resolve_dataset_path(raw) @@ -1578,9 +1575,8 @@ async def diffusion_training_info(current_subject: str = Depends(get_current_sub children = sorted( p for p in root.iterdir() - # Skip symlinked dirs: the CRUD resolver (_resolve_dataset_folder) rejects a - # symlinked dataset, so discovery must not advertise one as selectable (an external - # directory the read/caption/delete routes would then refuse). + # Skip symlinked dirs: the CRUD resolver rejects them, so discovery must not + # advertise one as selectable (the read/caption/delete routes would refuse it). if p.is_dir() and not p.is_symlink() and not p.name.startswith(".") ) except OSError: @@ -1760,11 +1756,9 @@ async def upload_diffusion_dataset( ), ) out.write(chunk) - # Reject a decompression bomb before commit: the byte-limit above passes a small, highly - # compressible PNG whose decoded pixels are huge, and the trainer later decodes every - # image in full when building its latent cache. Mirror the inference decode guard - # (diffusion._decode_b64_image) and bound each image's dimensions from the header, BEFORE - # any pixel decompression, so an oversized upload 400s here rather than OOMing the run. + # Reject a decompression bomb before commit: a small compressible PNG can pass the byte + # limit yet decode to huge pixels and OOM the trainer's latent cache, so bound each + # image's dimensions from the header (mirrors diffusion._decode_b64_image). if Path(filename).suffix.lower() in _DIFFUSION_DATASET_IMAGE_EXTS: _validate_uploaded_training_image(tmp, filename) uploaded += 1 @@ -1858,27 +1852,24 @@ def _resolve_dataset_folder(name: str, *, must_exist: bool = True) -> Path: return folder -# Bound each uploaded training image's dimensions (matches the inference decode guard's 4096px -# per-side limit in diffusion._decode_b64_image). A small, highly compressible PNG can smuggle huge -# pixel dimensions past the byte limit and OOM the trainer when it decodes the image for its latent -# cache, so reject an over-limit image from the header before any pixel decompression. +# Per-side dimension bound for uploaded training images, matching diffusion._decode_b64_image's +# 4096px inference guard, so a compressible PNG can't smuggle huge pixels past the byte limit. _MAX_TRAINING_IMAGE_SIDE = 4096 def _validate_uploaded_training_image(path: Path, original_name: str) -> None: """Reject an uploaded training image whose decoded dimensions exceed the per-side limit. - Reads only the image header (never img.load()), so a crafted small-payload / huge-dimension file - is caught before it can spike memory. Scoped to the decompression-bomb vector only: bytes PIL - cannot identify are left as-is (the upload contract accepts arbitrary bytes under an image - extension), so this changes behaviour solely for oversized real images.""" + Reads only the header (never img.load()), so a small-payload / huge-dimension file is caught + before it spikes memory. Bytes PIL cannot identify are left as-is (the upload contract accepts + arbitrary bytes under an image extension), so only oversized real images change behaviour.""" from PIL import Image, UnidentifiedImageError try: with Image.open(path) as image: width, height = image.size except (OSError, UnidentifiedImageError, ValueError): - return # not a decodable image -> not a decompression bomb; leave the existing contract + return # not a decodable image -> not a bomb; leave the existing contract if width > _MAX_TRAINING_IMAGE_SIDE or height > _MAX_TRAINING_IMAGE_SIDE: raise HTTPException( status_code = 400, @@ -1917,9 +1908,8 @@ def _load_metadata_captions(folder: Path) -> dict[str, str]: meta_path = folder / meta_name if not meta_path.is_file(): continue - # Tolerate a bad upload: invalid UTF-8 (UnicodeError, not an OSError), or a line that is - # valid JSON but not an object (``[]`` / ``null`` / a string / a number). Neither should - # 500 the info / labeling / caption / summary endpoints; the record is simply skipped. + # Tolerate a bad upload (invalid UTF-8, or a line of non-object JSON): skip the record so + # the info / labeling / caption / summary endpoints don't 500. try: lines = meta_path.read_text(encoding = "utf-8").splitlines() except (OSError, UnicodeError): diff --git a/studio/backend/routes/video.py b/studio/backend/routes/video.py index aed127fbad..b1ab370a45 100644 --- a/studio/backend/routes/video.py +++ b/studio/backend/routes/video.py @@ -133,11 +133,10 @@ async def load_video_model( ) if device != "cpu": - # Register the in-flight load UNDER the arbiter lock (not after acquire_for - # returns): a competing Images/chat acquire in that gap would otherwise evict - # VIDEO before begin_load marks a load in-flight, so eviction finds nothing to - # cancel and both loaders allocate VRAM at once. begin_load returns at once, so - # the lock is held only briefly. Mirrors the images/load handoff. + # Register the in-flight load UNDER the arbiter lock (not after acquire_for returns): + # otherwise a competing Images/chat acquire in that gap evicts VIDEO before the load is + # marked in-flight, finds nothing to cancel, and both loaders allocate VRAM at once. + # Mirrors the images/load handoff. status_dict = await asyncio.to_thread(acquire_for, VIDEO, _begin_load) else: await asyncio.to_thread(release, VIDEO) @@ -225,11 +224,9 @@ async def unload_video_model(current_subject: str = Depends(get_current_subject) backend = get_video_backend() status_dict = await asyncio.to_thread(backend.unload) # Drop VIDEO ownership only if nothing is resident AND no new load is in flight: a concurrent - # /video/load that re-acquired VIDEO must keep ownership (release is owner-guarded but - # identity-less, so an unconditional release would clear the newer claim). The idle check and - # the release must be ATOMIC (release_if): the load's acquire_for register runs under the same - # arbiter lock, so a plain check-then-release could pass the stale check and then clear the - # newer claim. Mirrors the images route. + # /video/load that re-acquired VIDEO must keep ownership. The idle check and release must be + # ATOMIC (release_if): the load's register runs under the same lock, so a plain + # check-then-release could clear the newer claim. Mirrors the images route. await asyncio.to_thread( release_if, VIDEO, diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index 4e1aa3a74d..ec455247ac 100644 --- a/studio/backend/tests/test_cached_gguf_routes.py +++ b/studio/backend/tests/test_cached_gguf_routes.py @@ -448,8 +448,7 @@ def test_list_cached_models_tags_diffusers_pipeline_as_text_to_image(monkeypatch [ _file("model_index.json", 1_000), _file("text_encoder/model.safetensors", 9_000), - # A complete pipeline carries its denoiser weights; without them the row is a - # companion-only prefetch and would be marked partial (see the dedicated test). + # A complete pipeline carries its denoiser weights; without them it would be partial. _file("transformer/diffusion_pytorch_model.safetensors", 9_000), ], tmp_path / "models--Tongyi-MAI--Z-Image-Turbo", @@ -475,11 +474,9 @@ def test_list_cached_models_tags_diffusers_pipeline_as_text_to_image(monkeypatch def test_list_cached_models_marks_companion_only_pipeline_partial(monkeypatch, tmp_path): - """A GGUF image load prefetches its companion base repo's VAE / text-encoder / model_index.json - but deliberately skips the multi-GB transformer (the GGUF supplies it). That snapshot carries a - root model_index.json yet is not a loadable BF16 pipeline, so it must be marked partial (the - picker drops partial rows) rather than advertised as fully on-device. A sibling repo that DOES - ship its transformer shards stays complete.""" + """A companion-only prefetch (VAE / text-encoder / model_index.json but no transformer) carries + a root model_index.json yet is not a loadable pipeline, so it must be marked partial. A sibling + repo that DOES ship its transformer shards stays complete.""" companion_only = _repo( "black-forest-labs/FLUX.1-dev", [ diff --git a/studio/backend/tests/test_gpu_arbiter.py b/studio/backend/tests/test_gpu_arbiter.py index 72f71aa6da..eabd4747cf 100644 --- a/studio/backend/tests/test_gpu_arbiter.py +++ b/studio/backend/tests/test_gpu_arbiter.py @@ -126,9 +126,8 @@ def test_release_if_by_non_owner_is_noop(calls): def test_release_if_predicate_sees_a_reregistered_same_owner_load(calls): - # The race release_if closes: a slow unload's predicate must observe a concurrent same-owner - # load that re-registered ownership, and NOT drop the newer claim. Simulate the re-register by - # having the predicate report a load now in flight; ownership must stay with DIFFUSION. + # The race release_if closes: a slow unload's predicate reports a load now in flight (a + # re-registered same-owner load), so ownership must stay with DIFFUSION. arb.acquire_for(arb.DIFFUSION) loading = {"in_flight": True} assert arb.release_if(arb.DIFFUSION, lambda: not loading["in_flight"]) is False @@ -165,9 +164,8 @@ def test_register_failure_leaves_ownership_in_place(calls): def test_competing_acquire_blocks_until_register_completes(monkeypatch): - # The window this closes: while DIFFUSION registers its load, a competing VIDEO acquire - # must not evict DIFFUSION until the load is marked in-flight. Holding the lock across - # register makes the competitor wait, so eviction never races an unregistered load. + # While DIFFUSION registers its load, a competing VIDEO acquire must block (not evict) until + # the load is in-flight; holding the lock across register makes eviction never race it. import threading import time diff --git a/studio/backend/tests/test_local_model_format.py b/studio/backend/tests/test_local_model_format.py index 75b8f20161..ebbc761215 100644 --- a/studio/backend/tests/test_local_model_format.py +++ b/studio/backend/tests/test_local_model_format.py @@ -118,11 +118,9 @@ def test_scan_models_dir_classifies_root_gguf_with_config(tmp_path): def test_scan_models_dir_surfaces_diffusers_pipeline_folder(tmp_path): - # A standard diffusers PIPELINE folder keeps its weights/configs in component subdirs - # (transformer/, vae/, ...) and carries only model_index.json at the root. The Images/Video - # load path accepts such a local pipeline dir, so the scan must surface it -- otherwise the - # weights-in-subdirs layout is missed and it never reaches task tagging / the On Device - # picker. It is not a GGUF, so model_format stays None (task tagging classifies it later). + # A diffusers PIPELINE folder (weights in component subdirs, only model_index.json at the root) + # is loadable, so the scan must surface it or it never reaches the On Device picker. Not a GGUF, + # so model_format stays None (task tagging classifies it later). root = tmp_path / "models" pipe = root / "my-pipeline" _touch(pipe / "model_index.json") @@ -137,10 +135,9 @@ def test_scan_models_dir_surfaces_diffusers_pipeline_folder(tmp_path): def test_scan_models_dir_surfaces_root_diffusers_pipeline(tmp_path): - # A custom scan folder can point DIRECTLY at a diffusers pipeline (not a parent of repos). - # Its weights live in component subdirs under a root model_index.json, so _is_model_directory - # rejects the root; without admitting it the scan would surface the component subdirs - # (transformer/, vae/) as bogus models and hide the real pipeline. Treat the root as one model. + # A scan folder can point DIRECTLY at a diffusers pipeline, which _is_model_directory rejects; + # without admitting it the scan surfaces the component subdirs as bogus models and hides the + # real pipeline. Treat the root as one model. root = tmp_path / "my-local-pipeline" _touch(root / "model_index.json") _touch(root / "transformer" / "config.json") @@ -188,10 +185,9 @@ def test_local_task_tags_family_named_pipeline_dir(tmp_path): def test_local_task_none_for_familyless_pipeline_dir(tmp_path): - # A generically named on-device pipeline (top-level model_index.json, no family token in its - # id / name / filename) is UNLOADABLE: the Images load path resolves no family via - # detect_family_for_pick and 400s after evicting the GPU owner. It must stay untagged so the - # picker never advertises a row that always fails; model_index.json alone is not enough. + # A generically named on-device pipeline (model_index.json, no family token) is UNLOADABLE: the + # Images load path resolves no family and 400s after evicting the GPU owner, so it must stay + # untagged and never be advertised. d = tmp_path / "my-local-pipeline" _touch(d / "model_index.json") _touch(d / "unet" / "diffusion_pytorch_model.safetensors") @@ -244,9 +240,8 @@ def test_local_task_tags_video_single_file_checkpoint(tmp_path): def test_local_task_tags_single_file_by_checkpoint_filename(tmp_path): - # A generically named folder holding one loadable checkpoint whose FILENAME identifies the - # family (the parent dir does not) is loadable -- the route resolves the sole file via - # resolve_local_single_file -- so tag it from the filename or the task-scoped picker hides it. + # A folder holding one checkpoint whose FILENAME identifies the family (not the parent dir) is + # loadable via resolve_local_single_file, so tag it from the filename or the picker hides it. d = tmp_path / "downloads" _touch(d / "qwen-image-2509.safetensors") # family only in the filename, no model_index.json m = _local(d, id = str(d), display_name = "downloads") diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index 28536294de..675520db0f 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -100,11 +100,9 @@ def test_loaded_repo_ids_includes_native_companions(): def test_loaded_repo_ids_tracks_variant_encoder_by_gguf_filename(): - # FLUX.2-klein-9B pairs with Qwen3-8B, and a local *klein-9B*.gguf carries that keyword only in - # the basename (not the repo id). loaded_repo_ids() must reproduce the committed load identity - # (repo id + GGUF filename), else it falls back to the 4B default encoder and the cache-deletion - # guard protects the wrong repo -- leaving the 8B encoder this load actually downloaded deletable - # while one-shot sd-cli still re-reads it every generation. + # A local *klein-9B*.gguf carries the variant keyword only in the basename, so loaded_repo_ids() + # must include the GGUF filename in the load identity; otherwise it falls back to the 4B default + # encoder and the cache-deletion guard protects the wrong repo. b = SdCppDiffusionBackend(engine = _FakeEngine()) fam = detect_family("flux.2-klein") b._state = bk._SdState( @@ -325,10 +323,9 @@ def test_generate_progress_tracks_parsed_steps(): def test_generate_publishes_progress_before_lora_resolution(monkeypatch): - # Native LoRA resolution (listing/downloading a not-yet-cached adapter) happens during the - # pre-generate setup while _generate_lock is already held. A reload/progress probe in that - # window must read ACTIVE, not idle, or the UI queues a second generate behind the first. - # So _gen is published before LoRA resolution, mirroring the diffusers path. + # LoRA resolution runs during pre-generate setup while _generate_lock is held, so a progress + # probe in that window must read ACTIVE; _gen is published before it, mirroring the diffusers + # path. from core.inference import diffusion_lora eng = _FakeEngine() diff --git a/studio/backend/tests/test_video_gallery.py b/studio/backend/tests/test_video_gallery.py index e6a3d6b361..fdc6f93f02 100644 --- a/studio/backend/tests/test_video_gallery.py +++ b/studio/backend/tests/test_video_gallery.py @@ -221,9 +221,8 @@ def test_list_skips_corrupt_sidecar(): def test_clear_preserves_mp4_with_present_but_invalid_sidecar(): - # A hand-dropped MP4 whose sidecar PARSES as JSON but lacks the required recipe keys (e.g. "{}") - # is hidden by list_videos (fails the GalleryVideo schema filter), so clear must not destroy it - # while removing the owned pair. Regression for the sidecar-validation gap. + # A hand-dropped MP4 whose sidecar parses but lacks the required recipe keys (e.g. "{}") is + # hidden by list_videos, so clear must not destroy it while removing the owned pair. directory = gallery.gallery_dir() (directory / "foreign.mp4").write_bytes(_mp4()) (directory / "foreign.json").write_text("{}", encoding = "utf-8") diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index ff6b814e55..6cf2d56fd8 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -1287,10 +1287,9 @@ export function ImagesPage({ active = true }: { active?: boolean }) { setNegativePrompt(image.guidance > 0 ? (image.negative_prompt ?? "") : ""); setSteps(image.steps); setGuidance(image.guidance); - // Restore from the BASE batch seed, not this image's own seed. The native engine derives - // per-image seeds as base + index, so replaying with the derived seed AND the original - // batch_size would advance a second time and reproduce a different image. Diffusers shares one - // seed, so batch_seed == seed there; older records without batch_seed fall back to seed. + // Restore from the BASE batch seed, not this image's own derived seed (base + index), or + // replaying with batch_size would advance again and reproduce a different image. Diffusers + // shares one seed (batch_seed == seed); older records without batch_seed fall back to seed. setSeed(String(image.batch_seed ?? image.seed)); setWidth(image.width); setHeight(image.height);