diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 0f92e3bba3..b3b1384aea 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -294,8 +294,8 @@ class DiffusionBackend: negative_prompt: Optional[str] = None, width: int = 1024, height: int = 1024, - steps: int = 24, - guidance: float = 3.5, + steps: int = 9, # Z-Image-Turbo: 9 steps = 8 DiT forwards (official default). + guidance: float = 0.0, # Turbo is distilled CFG-free; guidance must be 0. seed: Optional[int] = None, ) -> dict[str, Any]: import torch @@ -326,11 +326,9 @@ class DiffusionBackend: kwargs["negative_prompt"] = negative_prompt image = state.pipe(**kwargs).images[0] - return { - "image_b64": encode_png_base64(image), - "mime": "image/png", - "seed": int(seed), - } + # Return the PIL image (not yet encoded): the route embeds the + # generation recipe and persists it via the gallery. + return {"image": image, "seed": int(seed), "repo_id": state.repo_id} def unload(self) -> dict[str, Any]: with self._lock: diff --git a/studio/backend/core/inference/image_gallery.py b/studio/backend/core/inference/image_gallery.py new file mode 100644 index 0000000000..c495aa6e7d --- /dev/null +++ b/studio/backend/core/inference/image_gallery.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Disk-backed persistence for generated images. + +Each image is a PNG under ``studio_root()/images`` with its full generation +recipe embedded as PNG text chunks: a structured ``unsloth`` JSON blob (the +source of truth the gallery reads back) plus an Automatic1111-style +``parameters`` string for interop with other tools. Because the recipe lives +inside the file, a downloaded PNG carries its own settings. + +The gallery is intentionally dumb storage: the route owns the metadata schema +and passes a plain dict; this module only writes/reads/sorts files. +""" + +from __future__ import annotations + +import base64 +import json +import re +import uuid +from pathlib import Path +from typing import Any, Optional + +from loggers import get_logger +from utils.paths import ensure_dir, studio_root + +logger = get_logger(__name__) + +# PNG text-chunk key holding our structured recipe JSON. +_META_KEY = "unsloth" +# Image ids are file stems; restrict to filename-safe chars so a crafted id +# can't escape the gallery directory. +_ID_RE = re.compile(r"^[A-Za-z0-9_-]{1,128}$") + + +def gallery_dir() -> Path: + return ensure_dir(studio_root() / "images") + + +def _params_text(meta: dict[str, Any]) -> str: + """Automatic1111-style ``parameters`` string for cross-tool interop.""" + lines = [str(meta.get("prompt", ""))] + negative = meta.get("negative_prompt") + if negative: + lines.append(f"Negative prompt: {negative}") + lines.append( + f"Steps: {meta.get('steps')}, CFG scale: {meta.get('guidance')}, " + f"Seed: {meta.get('seed')}, Size: {meta.get('width')}x{meta.get('height')}, " + f"Model: {meta.get('model', '')}" + ) + return "\n".join(lines) + + +def _png_bytes(image: Any, meta: dict[str, Any]) -> bytes: + import io + + from PIL.PngImagePlugin import PngInfo + + info = PngInfo() + info.add_text(_META_KEY, json.dumps(meta)) + info.add_text("parameters", _params_text(meta)) + buf = io.BytesIO() + image.save(buf, format = "PNG", pnginfo = info) + return buf.getvalue() + + +def save(image: Any, meta: dict[str, Any]) -> tuple[dict[str, Any], str]: + """Persist a PIL image with its recipe embedded; return (record, base64 PNG). + + Returning the bytes we just wrote lets the caller hand them straight to the + client without reading the file back off disk.""" + image_id = uuid.uuid4().hex + png_bytes = _png_bytes(image, meta) + (gallery_dir() / f"{image_id}.png").write_bytes(png_bytes) + return _record(image_id, meta), base64.b64encode(png_bytes).decode("ascii") + + +def _record(image_id: str, meta: dict[str, Any]) -> dict[str, Any]: + return { + **meta, + "id": image_id, + "url": f"/api/inference/images/gallery/{image_id}/file", + } + + +def image_path(image_id: str) -> Optional[Path]: + """Resolve an id to its on-disk PNG, or None if missing / unsafe.""" + if not _ID_RE.match(image_id): + return None + path = gallery_dir() / f"{image_id}.png" + # Defence in depth: confirm the resolved path is still inside the gallery. + try: + path.resolve().relative_to(gallery_dir().resolve()) + except ValueError: + return None + return path if path.is_file() else None + + +def image_b64(image_id: str) -> Optional[str]: + path = image_path(image_id) + if path is None: + return None + return base64.b64encode(path.read_bytes()).decode("ascii") + + +def _read_meta(path: Path) -> Optional[dict[str, Any]]: + from PIL import Image + + try: + with Image.open(path) as im: + raw = im.text.get(_META_KEY) # type: ignore[attr-defined] + except Exception: + return None + if not raw: + return None + try: + meta = json.loads(raw) + except (ValueError, TypeError): + return None + 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 = [] + try: + paths = list(gallery_dir().glob("*.png")) + except OSError: + return [] + for path in paths: + 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 + + +def delete(image_id: str) -> bool: + path = image_path(image_id) + if path is None: + return False + try: + path.unlink() + return True + except OSError as exc: + logger.warning("image_gallery.delete_failed: %s", exc) + return False + + +def clear() -> int: + """Delete every gallery PNG; return how many were removed.""" + removed = 0 + try: + paths = list(gallery_dir().glob("*.png")) + except OSError: + return 0 + for path in paths: + try: + path.unlink() + removed += 1 + except OSError: + continue + return removed diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index ada2c5e526..a5326a382c 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1706,30 +1706,53 @@ class DiffusionGenerateRequest(BaseModel): prompt: str = Field(..., min_length = 1, description = "Text prompt") negative_prompt: Optional[str] = Field(None, description = "What to avoid (if the model supports it)") - width: int = Field(1024, ge = 256, le = 2048, description = "Image width in pixels (multiple of 8)") - height: int = Field(1024, ge = 256, le = 2048, description = "Image height in pixels (multiple of 8)") - steps: int = Field(24, ge = 1, le = 100, description = "Number of denoising steps") - guidance: float = Field(3.5, ge = 0.0, le = 20.0, description = "Classifier-free guidance scale") + width: int = Field(1024, ge = 256, le = 2048, description = "Image width in pixels (multiple of 16)") + height: int = Field(1024, ge = 256, le = 2048, description = "Image height in pixels (multiple of 16)") + steps: int = Field(9, ge = 1, le = 100, description = "Number of denoising steps") + guidance: float = Field(0.0, ge = 0.0, le = 20.0, description = "Classifier-free guidance scale") seed: Optional[int] = Field( None, ge = 0, le = 2**64 - 1, description = "Seed for reproducibility (random if omitted)" ) @field_validator("width", "height") @classmethod - def _multiple_of_8(cls, value: int) -> int: - # VAEs downsample by 8; non-multiples crash deep in the pipeline, so - # reject them here for a clean 422 instead of a cryptic 500. - if value % 8 != 0: - raise ValueError("must be a multiple of 8") + def _multiple_of_16(cls, value: int) -> int: + # Z-Image requires dimensions divisible by 16 (8x VAE downsample + 2x + # patch). Non-multiples crash deep in the pipeline, so reject them here + # for a clean 422 instead of a cryptic 500. + if value % 16 != 0: + raise ValueError("must be a multiple of 16") return value -class DiffusionGenerateResponse(BaseModel): - """A generated image plus the seed actually used.""" +class GalleryImage(BaseModel): + """A persisted image's full generation recipe (embedded in the PNG too).""" - image_b64: str = Field(..., description = "Base64-encoded PNG") + id: str = Field(..., description = "Stable id (the on-disk filename stem)") + url: str = Field(..., description = "Relative URL to fetch the PNG bytes") + prompt: str = Field(..., description = "Prompt used") + negative_prompt: Optional[str] = Field(None, description = "Negative prompt, if any") + width: int = Field(..., description = "Image width") + height: int = Field(..., description = "Image height") + steps: int = Field(..., description = "Denoising steps") + guidance: float = Field(..., description = "Guidance scale") + seed: int = Field(..., description = "Seed used") + model: Optional[str] = Field(None, description = "Model repo id that produced it") + created_at: float = Field(..., description = "Creation time (epoch seconds)") + + +class DiffusionGenerateResponse(BaseModel): + """A generated image (for instant display) plus its persisted gallery record.""" + + image_b64: str = Field(..., description = "Base64-encoded PNG (recipe embedded)") mime: str = Field("image/png", description = "MIME type of image_b64") - seed: int = Field(..., description = "Seed used (echoes the request, or the random one chosen)") + image: GalleryImage = Field(..., description = "The saved gallery record") + + +class GalleryListResponse(BaseModel): + """All persisted images, newest first.""" + + images: list[GalleryImage] = Field(default_factory = list) class DiffusionLoadProgressResponse(BaseModel): diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index c27f931636..0e06a7d4bc 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1054,6 +1054,8 @@ from models.inference import ( DiffusionGenerateResponse, DiffusionStatusResponse, DiffusionLoadProgressResponse, + GalleryImage, + GalleryListResponse, LoadResponse, LoadProgressResponse, UnloadResponse, @@ -10085,6 +10087,7 @@ async def generate_diffusion_image( request: DiffusionGenerateRequest, current_subject: str = Depends(get_current_subject), ): + from core.inference import image_gallery from core.inference.diffusion import get_diffusion_backend backend = get_diffusion_backend() @@ -10099,7 +10102,6 @@ async def generate_diffusion_image( guidance = request.guidance, seed = request.seed, ) - return DiffusionGenerateResponse(**result) except RuntimeError as exc: # No model loaded (or unloaded mid-flight) — a client-state problem. raise HTTPException(status_code = 409, detail = str(exc)) @@ -10107,6 +10109,79 @@ async def generate_diffusion_image( logger.error("diffusion.generate_failed: %s", exc) raise HTTPException(status_code = 500, detail = "Image generation failed.") + # Persist the image with its full recipe embedded, then hand the client both + # the bytes (instant display) and the saved record. + meta = { + "prompt": request.prompt, + "negative_prompt": request.negative_prompt, + "width": request.width, + "height": request.height, + "steps": request.steps, + "guidance": request.guidance, + "seed": result["seed"], + "model": result.get("repo_id"), + "created_at": time.time(), + } + try: + record, image_b64 = await asyncio.to_thread(image_gallery.save, result["image"], meta) + except Exception as exc: + logger.error("diffusion.persist_failed: %s", exc) + raise HTTPException(status_code = 500, detail = "Failed to save the generated image.") + + return DiffusionGenerateResponse( + image_b64 = image_b64, + mime = "image/png", + image = GalleryImage(**record), + ) + + +@studio_router.get("/images/gallery", response_model = GalleryListResponse) +async def list_gallery_images(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]) + + +@studio_router.get("/images/gallery/{image_id}/file") +async def get_gallery_image_file( + image_id: str, + current_subject: str = Depends(get_current_subject), +): + from core.inference import image_gallery + + path = await asyncio.to_thread(image_gallery.image_path, image_id) + if path is None: + raise HTTPException(status_code = 404, detail = "Image not found.") + data = await asyncio.to_thread(path.read_bytes) + # Immutable content (id is unique per image), so let the browser cache it. + return Response( + content = data, + media_type = "image/png", + headers = {"Cache-Control": "private, max-age=31536000, immutable"}, + ) + + +@studio_router.delete("/images/gallery/{image_id}") +async def delete_gallery_image( + image_id: str, + current_subject: str = Depends(get_current_subject), +): + from core.inference import image_gallery + + deleted = await asyncio.to_thread(image_gallery.delete, image_id) + if not deleted: + raise HTTPException(status_code = 404, detail = "Image not found.") + return {"deleted": True} + + +@studio_router.delete("/images/gallery") +async def clear_gallery_images(current_subject: str = Depends(get_current_subject)): + from core.inference import image_gallery + + removed = await asyncio.to_thread(image_gallery.clear) + return {"removed": removed} + @studio_router.post("/images/unload", response_model = DiffusionStatusResponse) async def unload_diffusion_model(current_subject: str = Depends(get_current_subject)): diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 289fa367b7..e72ec18d5f 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -191,9 +191,9 @@ def test_load_generate_unload_gguf(fake_runtime, tmp_path): assert "transformer" in _FakePipeline.last gen = backend.generate(prompt = "a sloth", width = 512, height = 512, steps = 4, guidance = 3.0) - assert gen["mime"] == "image/png" assert gen["seed"] == 4242 # random seed reported back - assert isinstance(gen["image_b64"], str) and gen["image_b64"] + assert gen["repo_id"] == str(tmp_path) # echoed so the route can record the model + assert gen["image"] is not None # PIL image handed to the route for persistence gen2 = backend.generate(prompt = "again", seed = 99) assert gen2["seed"] == 99 diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index 3232d9b98c..dda2edb96f 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -16,6 +16,7 @@ from fastapi.testclient import TestClient import core.inference.diffusion as diffusion_module import core.inference.gpu_arbiter as gpu_arbiter +import core.inference.image_gallery as gallery_module from auth.authentication import get_current_subject from routes.inference import studio_router @@ -53,7 +54,9 @@ class _FakeBackend: def generate(self, *, seed = None, **kwargs): if not self.loaded: raise RuntimeError("No diffusion model is loaded.") - return {"image_b64": "QUJD", "mime": "image/png", "seed": seed if seed is not None else 4242} + # The real backend returns the PIL image; the route persists it. The fake + # returns a sentinel object since image_gallery is stubbed in the fixture. + return {"image": object(), "seed": seed if seed is not None else 4242, "repo_id": "x/z-image"} def unload(self): self.loaded = False @@ -76,7 +79,7 @@ def _unloaded_status(): @pytest.fixture -def client(monkeypatch): +def client(monkeypatch, tmp_path): backend = _FakeBackend() monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) # Isolate from the real GPU arbiter: reset ownership and stub the evictors so @@ -85,6 +88,35 @@ def client(monkeypatch): monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.CHAT, lambda: None) monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.DIFFUSION, lambda: None) + # In-memory gallery backed by tmp files, so routes exercise persistence wiring + # without PIL/real disk under studio_root. + store: dict[str, dict] = {} + + def _save(image, meta): + image_id = f"img{len(store)}" + (tmp_path / f"{image_id}.png").write_bytes(b"PNG") + record = {**meta, "id": image_id, "url": f"/api/inference/images/gallery/{image_id}/file"} + store[image_id] = record + return record, "QUJD" # (record, base64 PNG) + + def _clear(): + n = len(store) + store.clear() + return n + + 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), + ) + monkeypatch.setattr( + gallery_module, "image_path", + lambda i: (tmp_path / f"{i}.png") if i in store else None, + ) + monkeypatch.setattr(gallery_module, "delete", lambda i: store.pop(i, None) is not None) + monkeypatch.setattr(gallery_module, "clear", _clear) + app = FastAPI() app.include_router(studio_router, prefix = "/api/inference") app.dependency_overrides[get_current_subject] = lambda: "test-user" @@ -106,17 +138,33 @@ def test_load_generate_status_unload_roundtrip(client): gen = client.post("/api/inference/images/generate", json = {"prompt": "a sloth", "seed": 7}) assert gen.status_code == 200 gbody = gen.json() - assert gbody["mime"] == "image/png" and gbody["seed"] == 7 and gbody["image_b64"] + assert gbody["mime"] == "image/png" and gbody["image_b64"] + # The persisted record carries the full recipe back. + img = gbody["image"] + assert img["seed"] == 7 and img["prompt"] == "a sloth" and img["id"] + + # The image is now listable, fetchable, and deletable. + listed = client.get("/api/inference/images/gallery").json()["images"] + assert [i["id"] for i in listed] == [img["id"]] + assert client.get(img["url"]).status_code == 200 + assert client.delete(img["url"].removesuffix("/file")).status_code == 200 + assert client.get("/api/inference/images/gallery").json()["images"] == [] unloaded = client.post("/api/inference/images/unload") assert unloaded.status_code == 200 and unloaded.json()["loaded"] is False assert client.get("/api/inference/images/status").json()["loaded"] is False -def test_generate_rejects_non_multiple_of_8(client): +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"}) - resp = client.post("/api/inference/images/generate", json = {"prompt": "p", "width": 1001}) - assert resp.status_code == 422 + # Odd, and a multiple of 8 that isn't a multiple of 16: both rejected, since + # Z-Image requires dimensions divisible by 16. + for bad in (1001, 1000): + resp = client.post("/api/inference/images/generate", json = {"prompt": "p", "width": bad}) + assert resp.status_code == 422, bad + # A multiple of 16 is accepted. + ok = client.post("/api/inference/images/generate", json = {"prompt": "p", "width": 1024}) + assert ok.status_code == 200 def test_load_requires_gguf_filename(client): diff --git a/studio/backend/tests/test_image_gallery.py b/studio/backend/tests/test_image_gallery.py new file mode 100644 index 0000000000..8228aedc1d --- /dev/null +++ b/studio/backend/tests/test_image_gallery.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for the disk-backed image gallery: PNG-embedded recipe round-trips, +listing order, safe id handling, and delete/clear.""" + +from __future__ import annotations + +import base64 +import io + +import pytest + +import core.inference.image_gallery as gallery + +PIL = pytest.importorskip("PIL") +from PIL import Image # noqa: E402 + + +@pytest.fixture(autouse = True) +def _tmp_gallery(monkeypatch, tmp_path): + # Point the gallery at a throwaway root instead of ~/.unsloth/studio. + monkeypatch.setattr(gallery, "studio_root", lambda: tmp_path) + + +def _img(color = (10, 20, 30)): + return Image.new("RGB", (16, 16), color) + + +def _meta(**over): + base = { + "prompt": "a sloth", + "negative_prompt": None, + "width": 1024, + "height": 1024, + "steps": 9, + "guidance": 0.0, + "seed": 7, + "model": "unsloth/Z-Image-Turbo-GGUF", + "created_at": 100.0, + } + base.update(over) + return base + + +def test_save_embeds_recipe_and_round_trips(): + record, _b64 = gallery.save(_img(), _meta()) + assert record["id"] and record["url"].endswith(f"{record['id']}/file") + + # The recipe is embedded in the PNG itself (portable), not just in a sidecar. + raw = base64.b64decode(gallery.image_b64(record["id"])) + with Image.open(io.BytesIO(raw)) as im: + assert im.text["unsloth"] + assert "Negative prompt" not in im.text["parameters"] # none given + assert "Steps: 9" in im.text["parameters"] + + listed = gallery.list_images() + assert len(listed) == 1 + assert listed[0]["prompt"] == "a sloth" and listed[0]["seed"] == 7 + + +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)) + assert [r["id"] for r in gallery.list_images()] == [new["id"], old["id"]] + + +def test_negative_prompt_recorded_in_parameters(): + record, _ = gallery.save(_img(), _meta(negative_prompt = "blurry")) + raw = base64.b64decode(gallery.image_b64(record["id"])) + with Image.open(io.BytesIO(raw)) as im: + assert "Negative prompt: blurry" in im.text["parameters"] + + +def test_delete_and_clear(): + a, _ = gallery.save(_img(), _meta(prompt = "a")) + gallery.save(_img(), _meta(prompt = "b")) + assert gallery.delete(a["id"]) is True + assert gallery.delete(a["id"]) is False # already gone + assert len(gallery.list_images()) == 1 + assert gallery.clear() == 1 + assert gallery.list_images() == [] + + +def test_image_path_rejects_unsafe_ids(): + # Traversal / bad chars never resolve to a path. + assert gallery.image_path("../../etc/passwd") is None + assert gallery.image_path("a/b") is None + assert gallery.image_path("missing") is None + + +def test_list_skips_foreign_pngs(tmp_path): + # A PNG without our recipe chunk (user dropped a file) is ignored. + foreign = gallery.gallery_dir() / "foreign.png" + _img().save(foreign, format = "PNG") + gallery.save(_img(), _meta(prompt = "ours")) + listed = gallery.list_images() + assert [r["prompt"] for r in listed] == ["ours"] diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index a7000038e8..960f6379c4 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -41,10 +41,25 @@ export interface DiffusionGenerateRequest { seed?: number; } +// A persisted image's full generation recipe (also embedded in the PNG). +export interface GalleryImage { + id: string; + url: string; + prompt: string; + negative_prompt: string | null; + width: number; + height: number; + steps: number; + guidance: number; + seed: number; + model: string | null; + created_at: number; +} + export interface DiffusionGenerateResponse { image_b64: string; mime: string; - seed: number; + image: GalleryImage; } async function parseJson(response: Response): Promise { @@ -87,3 +102,28 @@ export async function generateDiffusionImage( 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"), + ); + return images; +} + +export async function deleteGalleryImage(id: string): Promise { + const res = await authFetch(`/api/inference/images/gallery/${id}`, { method: "DELETE" }); + if (!res.ok) throw new Error(await readFastApiError(res)); +} + +export async function clearGallery(): Promise { + const res = await authFetch("/api/inference/images/gallery", { method: "DELETE" }); + if (!res.ok) throw new Error(await readFastApiError(res)); +} + +/** Fetch a gallery PNG (auth-protected, so it can't be a plain ) and + * wrap it in an object URL. Callers must revoke the URL when done. */ +export async function fetchGalleryObjectUrl(url: string): Promise { + const res = await authFetch(url); + if (!res.ok) throw new Error(await readFastApiError(res)); + return URL.createObjectURL(await res.blob()); +} diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 595f997fce..814fbf85d0 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -2,11 +2,22 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { useCallback, useEffect, useMemo, useRef, useState, type ReactNode } from "react"; -import { Download01Icon, ImageAdd02Icon } from "@hugeicons/core-free-icons"; +import { + ArrowReloadHorizontalIcon, + Delete02Icon, + Download01Icon, + ImageAdd02Icon, + InformationCircleIcon, +} from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; +import { + Popover, + PopoverContent, + PopoverTrigger, +} from "@/components/ui/popover"; import { Select, SelectContent, @@ -18,6 +29,7 @@ import { Slider } from "@/components/ui/slider"; import { Spinner } from "@/components/ui/spinner"; import { Textarea } from "@/components/ui/textarea"; import { SectionCard } from "@/components/section-card"; +import { InfoHint } from "@/components/ui/info-hint"; import { ModelSelector } from "@/components/assistant-ui/model-selector"; import type { ModelOption, @@ -31,9 +43,13 @@ import { toast } from "@/lib/toast"; import { type DiffusionLoadProgress, type DiffusionStatus, + type GalleryImage, + deleteGalleryImage, + fetchGalleryObjectUrl, generateDiffusionImage, getDiffusionLoadProgress, getDiffusionStatus, + getGallery, loadDiffusionModel, unloadDiffusionModel, } from "./api"; @@ -52,35 +68,35 @@ const MODELS: ModelOption[] = [ { id: MODEL.repo_id, name: MODEL.label, description: "Text-to-image · GGUF", isGguf: true }, ]; -// Keep at most this many generated images in the session gallery. -const MAX_GALLERY = 50; - +// Z-Image's official ~1-megapixel resolution buckets (the 1024 grid from the +// Tongyi-MAI demo app). All divisible by 16, the model's required step. const RESOLUTIONS: Array<{ label: string; w: number; h: number }> = [ - { label: "Square 1024", w: 1024, h: 1024 }, - { label: "Square 768", w: 768, h: 768 }, - { label: "Portrait 832×1216", w: 832, h: 1216 }, - { label: "Landscape 1216×832", w: 1216, h: 832 }, + { label: "1024 × 1024 (1:1)", w: 1024, h: 1024 }, + { label: "1152 × 896 (9:7)", w: 1152, h: 896 }, + { label: "896 × 1152 (7:9)", w: 896, h: 1152 }, + { label: "1152 × 864 (4:3)", w: 1152, h: 864 }, + { label: "864 × 1152 (3:4)", w: 864, h: 1152 }, + { label: "1248 × 832 (3:2)", w: 1248, h: 832 }, + { label: "832 × 1248 (2:3)", w: 832, h: 1248 }, + { label: "1280 × 720 (16:9)", w: 1280, h: 720 }, + { label: "720 × 1280 (9:16)", w: 720, h: 1280 }, + { label: "1344 × 576 (21:9)", w: 1344, h: 576 }, + { label: "576 × 1344 (9:21)", w: 576, h: 1344 }, ]; -interface ResultItem { - id: number; - src: string; - prompt: string; - width: number; - height: number; - steps: number; - guidance: number; - seed: number; -} - -// Generated images live here (module scope) so they survive the page unmounting -// when you switch tabs — the route is lazy and remounts otherwise empty. -const imageSession: { - results: ResultItem[]; - selectedId: number | null; - nextId: number; +// The gallery is persisted on the backend (durable across reloads); this module +// cache only holds the last-fetched records + their object/data URLs so a tab +// switch re-renders instantly without a refetch flash. Object URLs live for the +// app's lifetime (revoked only on delete), so they stay valid across remounts. +const galleryCache: { + images: GalleryImage[]; + selectedId: string | null; quant: string | null; -} = { results: [], selectedId: null, nextId: 0, quant: 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() }; function downloadImage(src: string, seed: number) { const link = document.createElement("a"); @@ -89,6 +105,10 @@ function downloadImage(src: string, seed: number) { link.click(); } +function formatTimestamp(epochSeconds: number): string { + return new Date(epochSeconds * 1000).toLocaleString(); +} + // The chat tab's model-load toast styling, reused verbatim so the diffusion // load toast is visually identical (persistent, progress bar, same chrome). const LOAD_TOAST_CLASSNAMES = { @@ -151,6 +171,7 @@ const IDLE_PROGRESS: DiffusionLoadProgress = { // label + standard Slider + number input, same classes. function SliderField({ label, + hint, value, min, max, @@ -158,6 +179,7 @@ function SliderField({ onChange, }: { label: string; + hint?: ReactNode; value: number; min: number; max: number; @@ -166,7 +188,10 @@ function SliderField({ }) { return (
- {label} + + {label} + {hint && {hint}} +
- +
+ + {hint && {hint}} +
{children}
); } +// One labeled row in the recipe popover. +function RecipeRow({ + label, + value, + wrap, + mono, +}: { + label: string; + value: string; + wrap?: boolean; + mono?: boolean; +}) { + return ( +
+ {label} + + {value} + +
+ ); +} + +// The full generation recipe for an image, with a one-click "restore to inputs". +function RecipePopover({ + image, + onRestore, +}: { + image: GalleryImage; + onRestore: (image: GalleryImage) => void; +}) { + return ( + + + + + +
+

Generation settings

+

{formatTimestamp(image.created_at)}

+
+
+ + {image.negative_prompt ? ( + + ) : null} + {image.model ? : null} + + + + +
+
+ +
+ + + ); +} + type Busy = "loading" | "unloading" | "generating" | null; export function ImagesPage() { - const [quant, setQuant] = useState(imageSession.quant); + const [quant, setQuant] = useState(galleryCache.quant); const [prompt, setPrompt] = useState( "a tiny ginger sloth coding in a sunlit treehouse, photorealistic", ); const [negativePrompt, setNegativePrompt] = useState(""); const [resolutionIdx, setResolutionIdx] = useState(0); - const [steps, setSteps] = useState(8); // Z-Image-Turbo is distilled to ~8 NFE. - const [guidance, setGuidance] = useState(1.0); + // Z-Image-Turbo official defaults: 9 steps (= 8 DiT forwards), guidance 0 + // (distilled CFG-free; a negative prompt is ignored at this guidance). + const [steps, setSteps] = useState(9); + const [guidance, setGuidance] = useState(0.0); const [seed, setSeed] = useState(""); const [busy, setBusy] = useState(null); const [status, setStatus] = useState(null); - const [results, setResults] = useState(() => imageSession.results); - const [selectedId, setSelectedId] = useState(() => imageSession.selectedId); - // Stable, ever-increasing ids: results are prepended, so an array index would - // re-key every existing image on each new generation. - const nextResultId = useRef(imageSession.nextId); + // 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 [selectedId, setSelectedId] = useState(() => galleryCache.selectedId); + const [srcById, setSrcById] = useState>(() => + Object.fromEntries(galleryCache.srcById), + ); const pollTimer = useRef | null>(null); // The persistent load toast's id, so each poll updates it in place (chat-style). const loadToastId = useRef(null); @@ -229,19 +340,80 @@ export function ImagesPage() { loadToastId.current = null; }, []); - // Persist the gallery to module scope so a tab switch doesn't drop it. + // Mirror to the module cache so a tab switch re-renders instantly. useEffect(() => { - imageSession.results = results; - imageSession.selectedId = selectedId; - imageSession.nextId = nextResultId.current; - imageSession.quant = quant; - }, [results, selectedId, quant]); + galleryCache.images = images; + galleryCache.selectedId = selectedId; + galleryCache.quant = quant; + }, [images, selectedId, quant]); const resolution = RESOLUTIONS[resolutionIdx]; const selected = useMemo( - () => results.find((r) => r.id === selectedId) ?? results[0] ?? null, - [results, selectedId], + () => images.find((i) => i.id === selectedId) ?? images[0] ?? null, + [images, selectedId], ); + const selectedSrc = selected ? srcById[selected.id] : undefined; + + // Fetch (once) the object URL for a record's PNG; cached across remounts. + const ensureSrc = useCallback(async (image: GalleryImage) => { + if (galleryCache.srcById.has(image.id) || galleryCache.inflight.has(image.id)) return; + galleryCache.inflight.add(image.id); + try { + const url = await fetchGalleryObjectUrl(image.url); + galleryCache.srcById.set(image.id, url); + setSrcById((prev) => ({ ...prev, [image.id]: url })); + } catch { + // Leave it without a src; the tile shows a placeholder. + } finally { + galleryCache.inflight.delete(image.id); + } + }, []); + + const loadGallery = useCallback(async () => { + try { + const records = await getGallery(); + galleryCache.images = records; + setImages(records); + records.forEach((image) => void ensureSrc(image)); + } catch { + // Best-effort: a failed gallery load shouldn't block the page. + } + }, [ensureSrc]); + + useEffect(() => { + void loadGallery(); + }, [loadGallery]); + + const handleDelete = useCallback(async (id: string) => { + try { + await deleteGalleryImage(id); + } catch (err) { + toast.error(err instanceof Error ? err.message : "Failed to delete image"); + return; + } + const url = galleryCache.srcById.get(id); + if (url?.startsWith("blob:")) URL.revokeObjectURL(url); + galleryCache.srcById.delete(id); + setSrcById((prev) => { + const next = { ...prev }; + delete next[id]; + return next; + }); + setImages((prev) => prev.filter((i) => i.id !== id)); + setSelectedId((cur) => (cur === id ? null : cur)); + }, []); + + // Load an image's recipe back into the form inputs. + const restoreSettings = useCallback((image: GalleryImage) => { + setPrompt(image.prompt); + setNegativePrompt(image.negative_prompt ?? ""); + setSteps(image.steps); + setGuidance(image.guidance); + setSeed(String(image.seed)); + const idx = RESOLUTIONS.findIndex((r) => r.w === image.width && r.h === image.height); + if (idx >= 0) setResolutionIdx(idx); + toast.success("Settings restored to inputs"); + }, []); const refreshStatus = useCallback(async () => { try { @@ -364,25 +536,13 @@ export function ImagesPage() { guidance, seed: parsedSeed, }); - const id = nextResultId.current++; - setResults((prev) => - // Cap the gallery so the base64 PNGs (held in module scope across tab - // switches) can't grow without bound over a long session. - [ - { - id, - src: `data:${res.mime};base64,${res.image_b64}`, - prompt: prompt.trim(), - width: resolution.w, - height: resolution.h, - steps, - guidance, - seed: res.seed, - }, - ...prev, - ].slice(0, MAX_GALLERY), - ); - setSelectedId(id); + // Display the returned bytes immediately (no refetch); the record is the + // durable, backend-persisted gallery entry. + const dataUrl = `data:${res.mime};base64,${res.image_b64}`; + galleryCache.srcById.set(res.image.id, dataUrl); + setSrcById((prev) => ({ ...prev, [res.image.id]: dataUrl })); + setImages((prev) => [res.image, ...prev]); + setSelectedId(res.image.id); } catch (err) { toast.error(err instanceof Error ? err.message : "Image generation failed"); } finally { @@ -419,7 +579,10 @@ export function ImagesPage() {