Persist the image gallery to disk with recipes embedded in each PNG, and match official Z-Image defaults
This commit is contained in:
oobabooga
parent
adebdfc9f4
commit
caa7efecc0
9 changed files with 753 additions and 111 deletions
|
|
@ -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:
|
||||
|
|
|
|||
165
studio/backend/core/inference/image_gallery.py
Normal file
165
studio/backend/core/inference/image_gallery.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
98
studio/backend/tests/test_image_gallery.py
Normal file
98
studio/backend/tests/test_image_gallery.py
Normal file
|
|
@ -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"]
|
||||
|
|
@ -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<T>(response: Response): Promise<T> {
|
||||
|
|
@ -87,3 +102,28 @@ export async function generateDiffusionImage(
|
|||
export async function unloadDiffusionModel(): Promise<DiffusionStatus> {
|
||||
return parseJson(await authFetch("/api/inference/images/unload", { method: "POST" }));
|
||||
}
|
||||
|
||||
export async function getGallery(): Promise<GalleryImage[]> {
|
||||
const { images } = await parseJson<{ images: GalleryImage[] }>(
|
||||
await authFetch("/api/inference/images/gallery"),
|
||||
);
|
||||
return images;
|
||||
}
|
||||
|
||||
export async function deleteGalleryImage(id: string): Promise<void> {
|
||||
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<void> {
|
||||
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 <img src>) and
|
||||
* wrap it in an object URL. Callers must revoke the URL when done. */
|
||||
export async function fetchGalleryObjectUrl(url: string): Promise<string> {
|
||||
const res = await authFetch(url);
|
||||
if (!res.ok) throw new Error(await readFastApiError(res));
|
||||
return URL.createObjectURL(await res.blob());
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<string, string>;
|
||||
// Ids with a fetch in flight, so concurrent ensureSrc calls don't double-fetch
|
||||
// (and leak the duplicate object URL).
|
||||
inflight: Set<string>;
|
||||
} = { 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 (
|
||||
<div className="flex items-center justify-between">
|
||||
<span className="text-xs font-medium text-muted-foreground">{label}</span>
|
||||
<span className="flex items-center gap-1 text-xs font-medium text-muted-foreground">
|
||||
{label}
|
||||
{hint && <InfoHint>{hint}</InfoHint>}
|
||||
</span>
|
||||
<div className="flex items-center gap-3">
|
||||
<Slider
|
||||
value={[value]}
|
||||
|
|
@ -191,35 +216,121 @@ function SliderField({
|
|||
}
|
||||
|
||||
// Matches the field-label style used across Studio (export/chat settings).
|
||||
function Field({ label, children }: { label: string; children: ReactNode }) {
|
||||
function Field({
|
||||
label,
|
||||
hint,
|
||||
children,
|
||||
}: {
|
||||
label: string;
|
||||
hint?: ReactNode;
|
||||
children: ReactNode;
|
||||
}) {
|
||||
return (
|
||||
<div className="flex flex-col gap-1.5">
|
||||
<label className="text-xs font-medium text-muted-foreground">{label}</label>
|
||||
<div className="flex items-center gap-1">
|
||||
<label className="text-xs font-medium text-muted-foreground">{label}</label>
|
||||
{hint && <InfoHint>{hint}</InfoHint>}
|
||||
</div>
|
||||
{children}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
// One labeled row in the recipe popover.
|
||||
function RecipeRow({
|
||||
label,
|
||||
value,
|
||||
wrap,
|
||||
mono,
|
||||
}: {
|
||||
label: string;
|
||||
value: string;
|
||||
wrap?: boolean;
|
||||
mono?: boolean;
|
||||
}) {
|
||||
return (
|
||||
<div className={cn("grid grid-cols-[72px_1fr] gap-2", wrap ? "items-start" : "items-center")}>
|
||||
<span className="text-muted-foreground">{label}</span>
|
||||
<span
|
||||
className={cn(
|
||||
"min-w-0 text-foreground",
|
||||
wrap ? "whitespace-pre-wrap break-words" : "truncate",
|
||||
mono && "font-mono",
|
||||
)}
|
||||
>
|
||||
{value}
|
||||
</span>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
// 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 (
|
||||
<Popover>
|
||||
<PopoverTrigger asChild>
|
||||
<Button size="sm" variant="secondary" className="gap-1.5">
|
||||
<HugeiconsIcon icon={InformationCircleIcon} className="size-4" />
|
||||
Recipe
|
||||
</Button>
|
||||
</PopoverTrigger>
|
||||
<PopoverContent align="end" side="top" className="w-80 p-0">
|
||||
<div className="border-b border-border/60 px-4 py-2.5">
|
||||
<p className="text-sm font-semibold">Generation settings</p>
|
||||
<p className="text-[11px] text-muted-foreground">{formatTimestamp(image.created_at)}</p>
|
||||
</div>
|
||||
<div className="flex flex-col gap-2 px-4 py-3 text-xs">
|
||||
<RecipeRow label="Prompt" value={image.prompt} wrap />
|
||||
{image.negative_prompt ? (
|
||||
<RecipeRow label="Negative" value={image.negative_prompt} wrap />
|
||||
) : null}
|
||||
{image.model ? <RecipeRow label="Model" value={image.model} /> : null}
|
||||
<RecipeRow label="Size" value={`${image.width} × ${image.height}`} />
|
||||
<RecipeRow label="Steps" value={String(image.steps)} />
|
||||
<RecipeRow label="Guidance" value={String(image.guidance)} />
|
||||
<RecipeRow label="Seed" value={String(image.seed)} mono />
|
||||
</div>
|
||||
<div className="border-t border-border/60 px-3 py-2.5">
|
||||
<Button size="sm" className="w-full gap-1.5" onClick={() => onRestore(image)}>
|
||||
<HugeiconsIcon icon={ArrowReloadHorizontalIcon} className="size-4" />
|
||||
Restore these settings
|
||||
</Button>
|
||||
</div>
|
||||
</PopoverContent>
|
||||
</Popover>
|
||||
);
|
||||
}
|
||||
|
||||
type Busy = "loading" | "unloading" | "generating" | null;
|
||||
|
||||
export function ImagesPage() {
|
||||
const [quant, setQuant] = useState<string | null>(imageSession.quant);
|
||||
const [quant, setQuant] = useState<string | null>(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<Busy>(null);
|
||||
const [status, setStatus] = useState<DiffusionStatus | null>(null);
|
||||
const [results, setResults] = useState<ResultItem[]>(() => imageSession.results);
|
||||
const [selectedId, setSelectedId] = useState<number | null>(() => 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<GalleryImage[]>(() => galleryCache.images);
|
||||
const [selectedId, setSelectedId] = useState<string | null>(() => galleryCache.selectedId);
|
||||
const [srcById, setSrcById] = useState<Record<string, string>>(() =>
|
||||
Object.fromEntries(galleryCache.srcById),
|
||||
);
|
||||
const pollTimer = useRef<ReturnType<typeof setTimeout> | null>(null);
|
||||
// The persistent load toast's id, so each poll updates it in place (chat-style).
|
||||
const loadToastId = useRef<string | number | null>(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() {
|
|||
<Field label="Prompt">
|
||||
<Textarea rows={4} value={prompt} onChange={(e) => setPrompt(e.target.value)} />
|
||||
</Field>
|
||||
<Field label="Negative prompt">
|
||||
<Field
|
||||
label="Negative prompt"
|
||||
hint="Z-Image-Turbo runs guidance-free, so a negative prompt is ignored — it only takes effect when guidance is above 0."
|
||||
>
|
||||
<Textarea
|
||||
rows={2}
|
||||
placeholder="What to avoid (optional)"
|
||||
|
|
@ -428,7 +591,10 @@ export function ImagesPage() {
|
|||
/>
|
||||
</Field>
|
||||
|
||||
<Field label="Resolution">
|
||||
<Field
|
||||
label="Resolution"
|
||||
hint="Z-Image's official ~1-megapixel resolutions. Every option sits on the 1024 grid the model was trained on; dimensions are multiples of 16."
|
||||
>
|
||||
<Select value={String(resolutionIdx)} onValueChange={(v) => setResolutionIdx(Number(v))}>
|
||||
<SelectTrigger>
|
||||
<SelectValue />
|
||||
|
|
@ -443,16 +609,25 @@ export function ImagesPage() {
|
|||
</Select>
|
||||
</Field>
|
||||
|
||||
<SliderField label="Steps" value={steps} min={1} max={50} step={1} onChange={setSteps} />
|
||||
<SliderField
|
||||
label="Steps"
|
||||
hint="Z-Image-Turbo is distilled to ~8 forward passes; 9 steps is the official setting. More steps rarely help."
|
||||
value={steps}
|
||||
min={1}
|
||||
max={50}
|
||||
step={1}
|
||||
onChange={setSteps}
|
||||
/>
|
||||
<SliderField
|
||||
label="Guidance"
|
||||
hint="Z-Image-Turbo is distilled CFG-free — keep this at 0. Higher values degrade Turbo output (other models use guidance)."
|
||||
value={guidance}
|
||||
min={0}
|
||||
max={15}
|
||||
step={0.5}
|
||||
onChange={setGuidance}
|
||||
/>
|
||||
<Field label="Seed">
|
||||
<Field label="Seed" hint="Leave empty for a fresh random seed each run.">
|
||||
<Input
|
||||
placeholder="Random if empty"
|
||||
value={seed}
|
||||
|
|
@ -468,10 +643,10 @@ export function ImagesPage() {
|
|||
|
||||
<div className="bg-card corner-squircle relative flex min-w-0 flex-1 flex-col overflow-hidden rounded-3xl ring-1 ring-foreground/10">
|
||||
<div className="relative flex flex-1 items-center justify-center overflow-auto p-6">
|
||||
{selected ? (
|
||||
{selected && selectedSrc ? (
|
||||
<>
|
||||
<img
|
||||
src={selected.src}
|
||||
src={selectedSrc}
|
||||
alt={selected.prompt}
|
||||
className="max-h-full max-w-full rounded-xl object-contain shadow-sm"
|
||||
/>
|
||||
|
|
@ -479,21 +654,31 @@ export function ImagesPage() {
|
|||
<span className="rounded-md bg-background/80 px-2 py-1 text-xs text-muted-foreground backdrop-blur">
|
||||
{selected.width}×{selected.height} · seed {selected.seed}
|
||||
</span>
|
||||
<RecipePopover image={selected} onRestore={restoreSettings} />
|
||||
<Button
|
||||
size="sm"
|
||||
variant="secondary"
|
||||
onClick={() => downloadImage(selected.src, selected.seed)}
|
||||
onClick={() => downloadImage(selectedSrc, selected.seed)}
|
||||
>
|
||||
<HugeiconsIcon icon={Download01Icon} className="mr-1.5 size-4" />
|
||||
Download
|
||||
</Button>
|
||||
<Button
|
||||
size="sm"
|
||||
variant="secondary"
|
||||
aria-label="Delete image"
|
||||
onClick={() => void handleDelete(selected.id)}
|
||||
>
|
||||
<HugeiconsIcon icon={Delete02Icon} className="size-4" />
|
||||
</Button>
|
||||
</div>
|
||||
</>
|
||||
) : busy === "generating" ? (
|
||||
// First image (no past result to keep showing): spin in place.
|
||||
) : busy === "generating" || selected ? (
|
||||
// First image generating, or the selected record's blob is still
|
||||
// loading — spin in place rather than flashing the empty state.
|
||||
<div className="flex flex-col items-center gap-3 text-muted-foreground">
|
||||
<Spinner className="size-8" />
|
||||
<p className="text-sm">Generating…</p>
|
||||
<p className="text-sm">{busy === "generating" ? "Generating…" : "Loading…"}</p>
|
||||
</div>
|
||||
) : (
|
||||
<div className="flex flex-col items-center gap-3 text-muted-foreground">
|
||||
|
|
@ -507,7 +692,7 @@ export function ImagesPage() {
|
|||
)}
|
||||
</div>
|
||||
|
||||
{(results.length > 0 || busy === "generating") && (
|
||||
{(images.length > 0 || busy === "generating") && (
|
||||
<div className="flex shrink-0 gap-2 overflow-x-auto border-t border-foreground/10 p-3">
|
||||
{/* In-progress generation: a placeholder tile at the front so past
|
||||
images stay visible and browsable while the new one renders. */}
|
||||
|
|
@ -516,18 +701,28 @@ export function ImagesPage() {
|
|||
<Spinner className="size-5 text-muted-foreground" />
|
||||
</div>
|
||||
)}
|
||||
{results.map((r) => (
|
||||
{images.map((image) => (
|
||||
<button
|
||||
key={r.id}
|
||||
key={image.id}
|
||||
type="button"
|
||||
onClick={() => setSelectedId(r.id)}
|
||||
title={`seed ${r.seed}`}
|
||||
className="relative size-16 shrink-0 overflow-hidden rounded-lg outline-none ring-1 ring-transparent transition-shadow hover:ring-border focus-visible:ring-2 focus-visible:ring-ring"
|
||||
onClick={() => setSelectedId(image.id)}
|
||||
title={`seed ${image.seed}`}
|
||||
className="relative size-16 shrink-0 overflow-hidden rounded-lg bg-muted/40 outline-none ring-1 ring-transparent transition-shadow hover:ring-border focus-visible:ring-2 focus-visible:ring-ring"
|
||||
>
|
||||
<img src={r.src} alt={r.prompt} className="size-full object-cover" />
|
||||
{srcById[image.id] ? (
|
||||
<img
|
||||
src={srcById[image.id]}
|
||||
alt={image.prompt}
|
||||
className="size-full object-cover"
|
||||
/>
|
||||
) : (
|
||||
<span className="flex size-full items-center justify-center">
|
||||
<Spinner className="size-4 text-muted-foreground" />
|
||||
</span>
|
||||
)}
|
||||
{/* Selection marker on a non-focusable overlay, so the button's
|
||||
own focus state can never mask it. */}
|
||||
{r.id === selected?.id && (
|
||||
{image.id === selected?.id && (
|
||||
<span className="pointer-events-none absolute inset-0 rounded-lg border-2 border-primary" />
|
||||
)}
|
||||
</button>
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue