diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index fa3ed47746..7e7ce2acfd 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -374,7 +374,7 @@ class DiffusionBackend: "active": True, "step": gen.step, "total_steps": gen.total_steps, - "fraction": min(gen.step / gen.total_steps, 1.0), + "fraction": gen.step / gen.total_steps, # step is 1..total, never over 1.0 "eta_seconds": gen.eta_seconds, } diff --git a/studio/backend/core/inference/image_gallery.py b/studio/backend/core/inference/image_gallery.py index 41e8e63cb2..4fb7990654 100644 --- a/studio/backend/core/inference/image_gallery.py +++ b/studio/backend/core/inference/image_gallery.py @@ -117,19 +117,31 @@ def _read_meta(path: Path) -> Optional[dict[str, Any]]: return meta if isinstance(meta, dict) else None -def list_images() -> list[dict[str, Any]]: - """All app-generated images, newest first (by embedded ``created_at``).""" - records = [] +def _mtime(path: Path) -> float: + try: + return path.stat().st_mtime + except OSError: + return 0.0 + + +def list_images(limit: Optional[int] = None, offset: int = 0) -> list[dict[str, Any]]: + """A newest-first window of images for infinite scroll. + + Ordered by file mtime (a cheap stat, ~= generation order) so a months-old + gallery isn't opened in full just to sort it; only the window's recipes are + read. limit=None returns everything from ``offset`` on.""" try: paths = list(gallery_dir().glob("*.png")) except OSError: return [] - for path in paths: + paths.sort(key = _mtime, reverse = True) + window = paths[offset:] if limit is None else paths[offset:offset + limit] + records = [] + for path in window: meta = _read_meta(path) if meta is None: # not one of ours (no recipe chunk) — skip continue records.append(_record(path.stem, meta)) - records.sort(key = lambda r: r.get("created_at", 0.0), reverse = True) return records diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 79539b8ae9..a492cc390c 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1754,9 +1754,10 @@ class DiffusionGenerateResponse(BaseModel): class GalleryListResponse(BaseModel): - """All persisted images, newest first.""" + """A newest-first page of persisted images, for infinite scroll.""" images: list[GalleryImage] = Field(default_factory = list) + has_more: bool = Field(False, description = "Whether older images remain past this page") class DiffusionGenerateProgressResponse(BaseModel): diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 6eed4a1d08..83ef7aad69 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10145,11 +10145,22 @@ async def generate_diffusion_image( @studio_router.get("/images/gallery", response_model = GalleryListResponse) -async def list_gallery_images(current_subject: str = Depends(get_current_subject)): +async def list_gallery_images( + limit: int = 50, + offset: int = 0, + current_subject: str = Depends(get_current_subject), +): from core.inference import image_gallery - records = await asyncio.to_thread(image_gallery.list_images) - return GalleryListResponse(images = [GalleryImage(**r) for r in records]) + limit = max(1, min(limit, 200)) + offset = max(0, offset) + # Fetch one extra to learn whether more remain, without a second scan. + records = await asyncio.to_thread(image_gallery.list_images, limit + 1, offset) + has_more = len(records) > limit + return GalleryListResponse( + images = [GalleryImage(**r) for r in records[:limit]], + has_more = has_more, + ) @studio_router.get("/images/gallery/{image_id}/file") diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 51ef8b246b..d01dce967a 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -308,18 +308,6 @@ def test_estimate_eta(): assert _estimate_eta(8, 8, first_step_at = 100.0, now = 107.0) == 0.0 -def test_generate_progress_reads_gen_state(): - from core.inference.diffusion import _GenState - - backend = DiffusionBackend() - assert backend.generate_progress()["active"] is False - - backend._gen = _GenState(total_steps = 8, step = 4, eta_seconds = 4.0) - p = backend.generate_progress() - assert p["active"] is True and p["step"] == 4 and p["total_steps"] == 8 - assert abs(p["fraction"] - 0.5) < 1e-9 and p["eta_seconds"] == 4.0 - - def test_begin_load_rejects_concurrent(monkeypatch): backend = DiffusionBackend() monkeypatch.setattr(DiffusionBackend, "_estimate_download_bytes", staticmethod(lambda *a, **k: 0)) diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index 0e941f5a77..1309746dbb 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -110,10 +110,11 @@ def client(monkeypatch, tmp_path): monkeypatch.setattr(gallery_module, "save", _save) monkeypatch.setattr(gallery_module, "image_b64", lambda i: "QUJD" if i in store else None) - monkeypatch.setattr( - gallery_module, "list_images", - lambda: sorted(store.values(), key = lambda r: r.get("created_at", 0.0), reverse = True), - ) + def _list_images(limit = None, offset = 0): + ordered = sorted(store.values(), key = lambda r: r.get("created_at", 0.0), reverse = True) + return ordered[offset:] if limit is None else ordered[offset:offset + limit] + + monkeypatch.setattr(gallery_module, "list_images", _list_images) monkeypatch.setattr( gallery_module, "image_path", lambda i: (tmp_path / f"{i}.png") if i in store else None, @@ -173,6 +174,15 @@ def test_generate_batch_size_persists_each_image(client): assert len(client.get("/api/inference/images/gallery").json()["images"]) == 3 +def test_gallery_pagination(client): + client.post("/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}) + client.post("/api/inference/images/generate", json = {"prompt": "p", "batch_size": 5, "seed": 1}) + page1 = client.get("/api/inference/images/gallery?limit=2&offset=0").json() + assert len(page1["images"]) == 2 and page1["has_more"] is True + last = client.get("/api/inference/images/gallery?limit=2&offset=4").json() + assert len(last["images"]) == 1 and last["has_more"] is False + + def test_generate_rejects_non_multiple_of_16(client): client.post("/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}) # Odd, and a multiple of 8 that isn't a multiple of 16: both rejected, since diff --git a/studio/backend/tests/test_image_gallery.py b/studio/backend/tests/test_image_gallery.py index 52e721dad6..b8a297c05e 100644 --- a/studio/backend/tests/test_image_gallery.py +++ b/studio/backend/tests/test_image_gallery.py @@ -8,6 +8,7 @@ from __future__ import annotations import base64 import io +import os import pytest @@ -59,12 +60,32 @@ def test_save_embeds_recipe_and_round_trips(): assert listed[0]["prompt"] == "a sloth" and listed[0]["seed"] == 7 +def _save_with_mtime(prompt: str, t: float) -> dict: + record = gallery.save(_img(), _meta(prompt = prompt, created_at = t)) + # Listing orders by mtime; set it explicitly so a tight test loop can't tie it. + os.utime(gallery.gallery_dir() / f"{record['id']}.png", (t, t)) + return record + + def test_list_is_newest_first(): - old = gallery.save(_img(), _meta(prompt = "old", created_at = 100.0)) - new = gallery.save(_img(), _meta(prompt = "new", created_at = 200.0)) + old = _save_with_mtime("old", 100.0) + new = _save_with_mtime("new", 200.0) assert [r["id"] for r in gallery.list_images()] == [new["id"], old["id"]] +def test_list_paginates_with_limit_offset(): + # 5 images, newest (t=4) first. + for i in range(5): + _save_with_mtime(f"p{i}", float(i)) + page1 = gallery.list_images(limit = 2, offset = 0) + page2 = gallery.list_images(limit = 2, offset = 2) + assert [r["prompt"] for r in page1] == ["p4", "p3"] + assert [r["prompt"] for r in page2] == ["p2", "p1"] + # limit=None still returns everything from the offset. + assert len(gallery.list_images()) == 5 + assert len(gallery.list_images(offset = 4)) == 1 + + def test_negative_prompt_recorded_in_parameters(): record = gallery.save(_img(), _meta(negative_prompt = "blurry")) raw = base64.b64decode(gallery.image_b64(record["id"])) diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index 6f2e6f623a..fb9236cb4b 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -115,11 +115,15 @@ export async function unloadDiffusionModel(): Promise { return parseJson(await authFetch("/api/inference/images/unload", { method: "POST" })); } -export async function getGallery(): Promise { - const { images } = await parseJson<{ images: GalleryImage[] }>( - await authFetch("/api/inference/images/gallery"), +export interface GalleryPage { + images: GalleryImage[]; + has_more: boolean; +} + +export async function getGallery(offset = 0, limit = 50): Promise { + return parseJson( + await authFetch(`/api/inference/images/gallery?offset=${offset}&limit=${limit}`), ); - return images; } export async function deleteGalleryImage(id: string): Promise { diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 0fef65b2cc..ab0e3b2723 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -106,13 +106,24 @@ function matchAspect(width: number, height: number): { key: string; portrait: bo // valid across remounts. const galleryCache: { images: GalleryImage[]; + hasMore: boolean; selectedId: string | null; quant: string | null; srcById: Map; // Ids with a fetch in flight, so concurrent ensureSrc calls don't double-fetch // (and leak the duplicate object URL). inflight: Set; -} = { images: [], selectedId: null, quant: null, srcById: new Map(), inflight: new Set() }; +} = { + images: [], + hasMore: false, + selectedId: null, + quant: null, + srcById: new Map(), + inflight: new Set(), +}; + +// Images loaded per infinite-scroll page. +const PAGE_SIZE = 50; // Export filename, e.g. Unsloth_20260624-143005_123.png. Batch siblings share // the seed + timestamp, so they get a "_" suffix past the first one. @@ -376,10 +387,13 @@ export function ImagesPage() { // Records come from the backend (durable); srcById maps each id to its object // URL (loaded images) or data URL (the one just generated). const [images, setImages] = useState(() => galleryCache.images); + const [hasMore, setHasMore] = useState(() => galleryCache.hasMore); const [selectedId, setSelectedId] = useState(() => galleryCache.selectedId); const [srcById, setSrcById] = useState>(() => Object.fromEntries(galleryCache.srcById), ); + // Guards a "load more" so a fast scroll can't fire several at once. + const loadingMore = useRef(false); const pollTimer = useRef | null>(null); // The persistent load toast's id, so each poll updates it in place (chat-style). const loadToastId = useRef(null); @@ -392,9 +406,10 @@ export function ImagesPage() { // Mirror to the module cache so a tab switch re-renders instantly. useEffect(() => { galleryCache.images = images; + galleryCache.hasMore = hasMore; galleryCache.selectedId = selectedId; galleryCache.quant = quant; - }, [images, selectedId, quant]); + }, [images, hasMore, selectedId, quant]); const selected = useMemo( () => images.find((i) => i.id === selectedId) ?? images[0] ?? null, @@ -419,15 +434,41 @@ export function ImagesPage() { const loadGallery = useCallback(async () => { try { - const records = await getGallery(); - galleryCache.images = records; - setImages(records); - records.forEach((image) => void ensureSrc(image)); + const page = await getGallery(0, PAGE_SIZE); + galleryCache.images = page.images; + galleryCache.hasMore = page.has_more; + setImages(page.images); + setHasMore(page.has_more); + page.images.forEach((image) => void ensureSrc(image)); } catch { // Best-effort: a failed gallery load shouldn't block the page. } }, [ensureSrc]); + // Load the next older page. offset = how many we've loaded so far. A newly + // generated image gets a newer mtime, so it sorts to the front on the backend + // too — the offset keeps pointing at the same older images. + const loadMore = useCallback(async () => { + if (loadingMore.current || !galleryCache.hasMore) return; + loadingMore.current = true; + try { + const page = await getGallery(galleryCache.images.length, PAGE_SIZE); + setImages((prev) => { + const seen = new Set(prev.map((i) => i.id)); + const next = [...prev, ...page.images.filter((i) => !seen.has(i.id))]; + galleryCache.images = next; + return next; + }); + galleryCache.hasMore = page.has_more; + setHasMore(page.has_more); + page.images.forEach((image) => void ensureSrc(image)); + } catch { + // transient; the user can scroll again to retry + } finally { + loadingMore.current = false; + } + }, [ensureSrc]); + useEffect(() => { void loadGallery(); }, [loadGallery]); @@ -872,7 +913,14 @@ export function ImagesPage() { {(images.length > 0 || busy === "generating") && ( -
+
{ + // Near the right edge: pull the next older page (infinite scroll). + const el = e.currentTarget; + if (el.scrollWidth - el.scrollLeft - el.clientWidth < 400) void loadMore(); + }} + > {/* In-progress generation: a placeholder tile at the front so past images stay visible and browsable while the new one renders. */} {busy === "generating" && ( @@ -906,6 +954,12 @@ export function ImagesPage() { )} ))} + {/* Tail spinner while older pages stream in on scroll. */} + {hasMore && ( +
+ +
+ )}
)}