diff --git a/studio/backend/core/training/diffusion_train_extras.py b/studio/backend/core/training/diffusion_train_extras.py index 8381d377d5..cf8cec4d95 100644 --- a/studio/backend/core/training/diffusion_train_extras.py +++ b/studio/backend/core/training/diffusion_train_extras.py @@ -171,6 +171,40 @@ def _sanitize(token: str) -> str: return re.sub(r"[^a-z0-9.]+", "-", str(token).lower()).strip("-") +def _hub_cache_roots() -> list[str]: + """Hub cache roots to look a repo up in, ACTIVE one first. + + Studio can move its cache during a session (Settings), and loading follows the live setting, + but ``huggingface_hub.constants.HF_HUB_CACHE`` is a snapshot of the environment at import + time. Reading only that constant left the revision "unresolved" (or pinned to a snapshot in + the previous root) once the cache moved, so pulling a new revision of the same checkpoint no + longer changed this key and a warm run silently reused the old embeddings and latents. The + constant stays as a fallback, since the trainer subprocess may run without Studio's settings + module importable.""" + import os # noqa: PLC0415 — keep the module import list light for the subprocess + + roots: list[str] = [] + try: + from utils.hf_cache_settings import active_hf_hub_cache # noqa: PLC0415 + + active = str(active_hf_hub_cache() or "").strip() + if active: + roots.append(active) + except Exception: # noqa: BLE001 -- settings unavailable in the subprocess: fall back + pass + for candidate in (os.environ.get("HF_HUB_CACHE"), os.environ.get("HUGGINGFACE_HUB_CACHE")): + if candidate and candidate.strip() and candidate.strip() not in roots: + roots.append(candidate.strip()) + try: + from huggingface_hub import constants # noqa: PLC0415 + + if constants.HF_HUB_CACHE and str(constants.HF_HUB_CACHE) not in roots: + roots.append(str(constants.HF_HUB_CACHE)) + except Exception: # noqa: BLE001 -- no hub package: whatever we collected above stands + pass + return roots + + def source_revision(ref: Any) -> str: """Revision/content marker for a checkpoint reference, resolved without loading it. @@ -208,21 +242,20 @@ def source_revision(ref: Any) -> str: ) return f"dir-{hashlib.sha256('|'.join(sorted(parts)).encode()).hexdigest()[:16]}" if "/" in name: - from huggingface_hub import constants # noqa: PLC0415 - org, _, repo = name.partition("/") - base = os.path.join(constants.HF_HUB_CACHE, f"models--{org}--{repo}") - ref_file = os.path.join(base, "refs", "main") - if os.path.isfile(ref_file): - with open(ref_file, encoding = "utf-8") as fh: - sha = fh.read().strip() - if sha: - return f"rev-{sha[:16]}" - snaps = os.path.join(base, "snapshots") - if os.path.isdir(snaps): - names = sorted(os.listdir(snaps)) - if len(names) == 1: - return f"rev-{names[0][:16]}" + for root in _hub_cache_roots(): + base = os.path.join(root, f"models--{org}--{repo}") + ref_file = os.path.join(base, "refs", "main") + if os.path.isfile(ref_file): + with open(ref_file, encoding = "utf-8") as fh: + sha = fh.read().strip() + if sha: + return f"rev-{sha[:16]}" + snaps = os.path.join(base, "snapshots") + if os.path.isdir(snaps): + names = sorted(os.listdir(snaps)) + if len(names) == 1: + return f"rev-{names[0][:16]}" return "unresolved" except Exception: # noqa: BLE001 — best-effort, never block a run return "unresolved" diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index b5e3208c9c..bedac02cac 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -6,6 +6,7 @@ Training API routes """ import sys +import threading from pathlib import Path from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile from fastapi.responses import StreamingResponse @@ -1580,6 +1581,38 @@ def _resolve_dataset_caption( return caption +_DATASET_IMPORT_LOCKS: Dict[str, "threading.Lock"] = {} +_DATASET_IMPORT_LOCKS_GUARD = threading.Lock() + + +def _dataset_import_lock(folder: Path) -> "threading.Lock": + """One lock per dataset folder, so two imports cannot fill the same empty name at once. + + Keyed by the resolved path (the same folder can be reached by different names), and kept for + the process lifetime: there are a handful of dataset folders and a Lock is tiny, while + dropping one while another thread holds it would defeat the point.""" + key = str(folder.resolve(strict = False)) + with _DATASET_IMPORT_LOCKS_GUARD: + lock = _DATASET_IMPORT_LOCKS.get(key) + if lock is None: + lock = threading.Lock() + _DATASET_IMPORT_LOCKS[key] = lock + return lock + + +def _import_response(entry: dict, folder: Path, *, imported: int) -> "DiffusionDatasetImportResponse": + summary = _diffusion_dataset_summary(folder) + return DiffusionDatasetImportResponse( + name = folder.name, + path = str(folder), + image_count = summary.image_count, + caption_count = summary.caption_count, + imported = imported, + license = entry["license"], + source_repo = entry["repo"], + ) + + def _diffusion_dataset_summary(folder: Path) -> DiffusionDatasetSummary: # Count an image as captioned only when it resolves to a NON-EMPTY caption via the same sidecar # over metadata precedence the trainer uses: an empty tombstone sidecar shadows a metadata row and @@ -2448,13 +2481,38 @@ async def import_diffusion_dataset_example( folder = _resolve_dataset_folder(body.name or entry["id"], must_exist = False) def do_import() -> DiffusionDatasetImportResponse: + folder.mkdir(parents = True, exist_ok = True) + if _diffusion_dataset_summary(folder).image_count > 0: + return _import_response(entry, folder, imported = 0) + # One import at a time per dataset folder. The training interlock COUNTS mutations rather + # than excluding them, so two imports of different examples into the same empty name both + # passed the emptiness check; the loser then merged its files into the winner's folder, + # overwriting same-numbered images and leaving a dataset whose images and captions came + # from two sources. Refusing the second is honest: the first is already filling that name. + lock = _dataset_import_lock(folder) + if not lock.acquire(blocking = False): + raise HTTPException( + status_code = 409, + detail = ( + f"An import into '{folder.name}' is already running. Wait for it to finish, " + "then reload the dataset list." + ), + ) + try: + return _do_import_locked(entry, folder) + finally: + lock.release() + + def _do_import_locked(entry: dict, folder: Path) -> DiffusionDatasetImportResponse: import os import shutil import tempfile - folder.mkdir(parents = True, exist_ok = True) - existing = _diffusion_dataset_summary(folder) imported = 0 + # Re-read under the lock: a winner may have promoted its staging dir while this request + # was checking, so the folder may no longer be empty. Returning it as-is matches the + # idempotent path rather than mixing two imports. + existing = _diffusion_dataset_summary(folder) if existing.image_count == 0: cap = int(entry["image_cap"]) # Materialize into a private staging dir and promote into the dataset folder only after the whole @@ -2495,15 +2553,6 @@ async def import_diffusion_dataset_example( os.replace(str(staging), str(folder)) finally: shutil.rmtree(staging, ignore_errors = True) - summary = _diffusion_dataset_summary(folder) - return DiffusionDatasetImportResponse( - name = folder.name, - path = str(folder), - image_count = summary.image_count, - caption_count = summary.caption_count, - imported = imported, - license = entry["license"], - source_repo = entry["repo"], - ) + return _import_response(entry, folder, imported = imported) return await asyncio.to_thread(do_import) diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py index 2361815f77..960d5dea20 100644 --- a/studio/backend/tests/test_diffusion_dataset_api.py +++ b/studio/backend/tests/test_diffusion_dataset_api.py @@ -729,3 +729,57 @@ def test_import_preserves_unrelated_files_when_folder_not_empty(client, ds_root, assert r.json()["imported"] == 3 assert sorted(p.name for p in folder.glob("*.png")) == [f"img_{i:04d}.png" for i in range(3)] assert (folder / "notes.md").read_text(encoding = "utf-8") == "keep me" + + +def test_a_second_concurrent_import_of_the_same_name_is_refused(client, ds_root, monkeypatch): + """The training interlock counts mutations rather than excluding them, so two imports into the + same empty name both got through; the loser then merged its files into the winner's folder and + produced a dataset built from two sources. The second request is refused instead.""" + import routes.training as training_route + + _install_fake_load_dataset(monkeypatch, n_rows = 3) + folder = ds_root / "my-tux" + folder.mkdir(parents = True) + # Stand in for the in-flight import: the lock is held for the whole materialize + promote. + held = training_route._dataset_import_lock(folder) + assert held.acquire(blocking = False) + try: + r = client.post( + "/api/train/diffusion/dataset/import-example", + json = {"id": "tuxemon", "name": "my-tux"}, + ) + finally: + held.release() + assert r.status_code == 409, r.text + assert "already running" in r.json()["detail"] + # Nothing was written into the folder the other import owns. + assert list(folder.glob("*.png")) == [] + + # Once it is free the same request imports normally. + ok = client.post( + "/api/train/diffusion/dataset/import-example", + json = {"id": "tuxemon", "name": "my-tux"}, + ) + assert ok.status_code == 200, ok.text + assert ok.json()["imported"] == 3 + + +def test_an_import_into_a_folder_filled_meanwhile_does_not_merge(client, ds_root, monkeypatch): + """The re-check under the lock: a request that waited while another import promoted its + staging dir must return the folder as-is, not move its own files on top.""" + import routes.training as training_route + + _install_fake_load_dataset(monkeypatch, n_rows = 3) + folder = ds_root / "my-tux" + folder.mkdir(parents = True) + (folder / "img_0000.png").write_bytes(b"\x89PNG\r\n\x1a\n" + b"0" * 16) + (folder / "img_0000.txt").write_text("from the first import", encoding = "utf-8") + + r = client.post( + "/api/train/diffusion/dataset/import-example", + json = {"id": "tuxemon", "name": "my-tux"}, + ) + assert r.status_code == 200, r.text + assert r.json()["imported"] == 0 + assert (folder / "img_0000.txt").read_text(encoding = "utf-8") == "from the first import" + assert not training_route._dataset_import_lock(folder).locked() diff --git a/studio/backend/tests/test_diffusion_train_extras.py b/studio/backend/tests/test_diffusion_train_extras.py index 550d63783d..428069d43f 100644 --- a/studio/backend/tests/test_diffusion_train_extras.py +++ b/studio/backend/tests/test_diffusion_train_extras.py @@ -292,3 +292,59 @@ def test_source_revision_marks_a_dir_update_and_never_raises(tmp_path): assert source_revision(str(d)) != third for ref in (None, "", "no/such/repo-xyz", "/does/not/exist", 7): assert isinstance(source_revision(ref), str) + + +def test_source_revision_reads_the_active_hub_cache(tmp_path, monkeypatch): + """Studio can move its HF cache mid-session and loading follows the live setting, but + huggingface_hub's HF_HUB_CACHE constant is a snapshot from import time. Reading only that left + the marker unresolved (or pinned to the old root), so pulling a new revision of the same + checkpoint stopped invalidating the conditioning cache and a warm run reused stale latents.""" + from core.training import diffusion_train_extras as extras + + old_root = tmp_path / "old" / "hub" + new_root = tmp_path / "new" / "hub" + for root, sha in ((old_root, "a" * 40), (new_root, "b" * 40)): + refs = root / "models--org--ckpt" / "refs" + refs.mkdir(parents = True) + (refs / "main").write_text(sha, encoding = "utf-8") + + monkeypatch.setattr(extras, "_hub_cache_roots", lambda: [str(new_root), str(old_root)]) + assert extras.source_revision("org/ckpt") == f"rev-{'b' * 16}" + + monkeypatch.setattr(extras, "_hub_cache_roots", lambda: [str(old_root)]) + assert extras.source_revision("org/ckpt") == f"rev-{'a' * 16}" + + +def test_hub_cache_roots_puts_the_active_studio_cache_first(monkeypatch, tmp_path): + from core.training import diffusion_train_extras as extras + from utils import hf_cache_settings + + active = tmp_path / "studio" / "hub" + monkeypatch.setattr(hf_cache_settings, "active_hf_hub_cache", lambda: str(active)) + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path / "env" / "hub")) + + roots = extras._hub_cache_roots() + assert roots and roots[0] == str(active) + # The environment root is still consulted, just after the live setting. + assert str(tmp_path / "env" / "hub") in roots + + +def test_hub_cache_roots_survives_without_studio_settings(monkeypatch, tmp_path): + # The trainer subprocess may run without Studio's settings module importable; the env and the + # library constant still have to work. + import builtins + + from core.training import diffusion_train_extras as extras + + real_import = builtins.__import__ + + def _blocked(name, *args, **kwargs): + if name == "utils.hf_cache_settings": + raise ImportError("no studio settings here") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", _blocked) + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path / "env" / "hub")) + roots = extras._hub_cache_roots() + monkeypatch.setattr(builtins, "__import__", real_import) + assert str(tmp_path / "env" / "hub") in roots diff --git a/studio/backend/tests/test_sd_cpp_install.py b/studio/backend/tests/test_sd_cpp_install.py index 1b1f32799e..deb43268e9 100644 --- a/studio/backend/tests/test_sd_cpp_install.py +++ b/studio/backend/tests/test_sd_cpp_install.py @@ -116,6 +116,22 @@ def test_windows_vulkan_picks_vulkan(): assert _resolve("Windows", "AMD64", "vulkan") == "sd-master-8caa3f9-bin-win-vulkan-x64.zip" +def test_windows_arm64_takes_no_x64_build(): + """An x64 zip cannot run natively on a Windows arm64 host. Reporting no match lets the caller + fall back instead of downloading and installing a binary that fails on first launch.""" + for accel in ("auto", "cuda", "vulkan", "rocm"): + assert _resolve("Windows", "ARM64", accel) is None + assert _resolve("Windows", "aarch64", accel) is None + + +def test_windows_arm64_picks_an_arm64_build_when_one_exists(): + assets = [*_ASSETS, "sd-master-8caa3f9-bin-win-avx2-arm64.zip"] + assert ( + resolve_release_asset(assets, system = "Windows", machine = "ARM64", accelerator = "auto") + == "sd-master-8caa3f9-bin-win-avx2-arm64.zip" + ) + + # ── cudart helper archive is never chosen as the engine ───────────────────── diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index bf31b685e7..5f827bbf2c 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -290,7 +290,18 @@ export async function generateDiffusionImage( ); } if (!response.ok && RESPONSE_LOST_STATUSES.has(response.status)) { - throw new GenerateResponseLostError(await readFastApiError(response)); + const detail = await readFastApiError(response); + // A proxy answers with HTML (or nothing); the app answers with a JSON body. Only the former + // means the request may still be running -- settling an application error would poll for a + // generation that never started and then report "did not reach the server" instead of the + // reason the backend gave. + if ( + response.status === 503 && + (response.headers.get("content-type") || "").toLowerCase().includes("application/json") + ) { + throw new Error(detail); + } + throw new GenerateResponseLostError(detail); } return parseJson(response); } diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index afcc54c282..98996250ec 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -1197,6 +1197,9 @@ export function ImagesPage({ active = true }: { active?: boolean }) { ); // Guards a "load more" so a fast scroll can't fire several at once. const loadingMore = useRef(false); + // The gallery strip, used as the IntersectionObserver root so a tile's PNG is fetched as it + // nears view instead of every tile of every page up front. + const stripRef = useRef(null); // False once the page truly unmounts (app close / chat-only eject). The page now // stays mounted across tab switches, so a switch does NOT flip this -- a batch keeps // generating off-tab; the multi-run loop only stops on a real unmount. @@ -1390,7 +1393,11 @@ export function ImagesPage({ active = true }: { active?: boolean }) { galleryCache.hasMore = page.has_more; setImages(page.images); setHasMore(page.has_more); - page.images.forEach((image) => void ensureSrc(image)); + // No visibility signal without IntersectionObserver (jsdom / an old webview), so keep the + // eager fetch there rather than render a strip of permanent spinners. + if (typeof IntersectionObserver === "undefined") { + page.images.forEach((image) => void ensureSrc(image)); + } } catch { // Best-effort: a failed gallery load shouldn't block the page. } @@ -1412,7 +1419,9 @@ export function ImagesPage({ active = true }: { active?: boolean }) { }); galleryCache.hasMore = page.has_more; setHasMore(page.has_more); - page.images.forEach((image) => void ensureSrc(image)); + if (typeof IntersectionObserver === "undefined") { + page.images.forEach((image) => void ensureSrc(image)); + } } catch { // transient; the user can scroll again to retry } finally { @@ -1420,6 +1429,43 @@ export function ImagesPage({ active = true }: { active?: boolean }) { } }, [ensureSrc]); + // A gallery page holds PAGE_SIZE multi-megabyte PNGs, and an object URL lives until the page + // closes, so fetching every record of every page up front grew memory without bound as the user + // scrolled -- for tiles they may never look at. Fetch a tile as it nears the strip's edge + // instead; it already shows a spinner until its src lands. Observed from here rather than + // through a ref on each tile, and re-run per page so "load more" tiles are picked up and + // removed ones go with the observer. Mirrors the Video page's strip. + useEffect(() => { + const root = stripRef.current; + if (!root || typeof IntersectionObserver === "undefined") return; + const io = new IntersectionObserver( + (entries) => { + for (const entry of entries) { + if (!entry.isIntersecting) continue; + const id = (entry.target as HTMLElement).dataset.imageId; + const image = id ? images.find((i) => i.id === id) : undefined; + if (image) void ensureSrc(image); + } + }, + // rootMargin applies to the ROOT's box only, so the root must be the horizontally + // scrolling strip: a tile past its right edge is clipped by the strip, and a margin on the + // viewport root would never reach it. The sideways margin starts a fetch a few tiles early. + { root, rootMargin: "0px 600px" }, + ); + for (const tile of root.querySelectorAll("[data-image-id]")) io.observe(tile); + return () => io.disconnect(); + }, [images, ensureSrc]); + + // The preview is what the user actually looks at, so the selected image is fetched whether or + // not its tile is on screen: a selection restored across a tab switch, a freshly generated + // image, or a pick made just before the strip scrolls it out of view. + useEffect(() => { + if (!selected) return; + void (async () => { + await ensureSrc(selected); + })(); + }, [selected, ensureSrc]); + useEffect(() => { void loadGallery(); }, [loadGallery]); @@ -3361,6 +3407,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { {(images.length > 0 || busy === "generating") && (
{ // Near the right edge: pull the next older page (infinite scroll). @@ -3379,6 +3426,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {