From eced3620fa2e84cf8806b9eeaa138f6ee3a5c6b8 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Wed, 24 Jun 2026 18:02:38 -0300 Subject: [PATCH] Add batch size and sequential count to image generation with reproducible per-run seeds --- studio/backend/core/inference/diffusion.py | 12 ++- .../backend/core/inference/image_gallery.py | 12 +-- studio/backend/models/inference.py | 11 ++- studio/backend/routes/inference.py | 43 ++++++----- .../backend/tests/test_diffusion_backend.py | 9 ++- studio/backend/tests/test_diffusion_routes.py | 36 ++++++--- studio/backend/tests/test_image_gallery.py | 10 +-- studio/frontend/src/features/images/api.ts | 5 +- .../src/features/images/images-page.tsx | 76 ++++++++++++++----- 9 files changed, 139 insertions(+), 75 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index b3b1384aea..83a36ecb0b 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -297,6 +297,7 @@ class DiffusionBackend: 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, + batch_size: int = 1, ) -> dict[str, Any]: import torch @@ -321,14 +322,17 @@ class DiffusionBackend: "num_inference_steps": steps, "guidance_scale": guidance, "generator": generator, + # Generate the whole batch in one forward pass (VRAM-heavy). All + # share this call's seed, drawn sequentially from one generator. + "num_images_per_prompt": batch_size, } if negative_prompt: kwargs["negative_prompt"] = negative_prompt - image = state.pipe(**kwargs).images[0] - # 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} + images = state.pipe(**kwargs).images + # Return the PIL images (not yet encoded): the route embeds each + # image's recipe and persists it via the gallery. + return {"images": list(images), "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 index c495aa6e7d..41e8e63cb2 100644 --- a/studio/backend/core/inference/image_gallery.py +++ b/studio/backend/core/inference/image_gallery.py @@ -65,15 +65,11 @@ def _png_bytes(image: Any, meta: dict[str, Any]) -> bytes: 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.""" +def save(image: Any, meta: dict[str, Any]) -> dict[str, Any]: + """Persist a PIL image with its recipe embedded; return the gallery record.""" 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") + (gallery_dir() / f"{image_id}.png").write_bytes(_png_bytes(image, meta)) + return _record(image_id, meta) def _record(image_id: str, meta: dict[str, Any]) -> dict[str, Any]: diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index a5326a382c..0a5e18758a 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1713,6 +1713,9 @@ class DiffusionGenerateRequest(BaseModel): seed: Optional[int] = Field( None, ge = 0, le = 2**64 - 1, description = "Seed for reproducibility (random if omitted)" ) + batch_size: int = Field( + 1, ge = 1, le = 32, description = "Images generated in one forward pass (VRAM-heavy)" + ) @field_validator("width", "height") @classmethod @@ -1742,11 +1745,11 @@ class GalleryImage(BaseModel): class DiffusionGenerateResponse(BaseModel): - """A generated image (for instant display) plus its persisted gallery record.""" + """The persisted gallery records for one generation call (a batch).""" - image_b64: str = Field(..., description = "Base64-encoded PNG (recipe embedded)") - mime: str = Field("image/png", description = "MIME type of image_b64") - image: GalleryImage = Field(..., description = "The saved gallery record") + images: list[GalleryImage] = Field( + ..., description = "Saved records, one per image in the batch" + ) class GalleryListResponse(BaseModel): diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 0e06a7d4bc..96c33f8528 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10101,6 +10101,7 @@ async def generate_diffusion_image( steps = request.steps, guidance = request.guidance, seed = request.seed, + batch_size = request.batch_size, ) except RuntimeError as exc: # No model loaded (or unloaded mid-flight) — a client-state problem. @@ -10109,30 +10110,34 @@ 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(), - } + # Persist each image with its full recipe embedded. The whole batch shares + # one seed (drawn sequentially from one generator) and one timestamp (the + # images are generated together), so their gallery order is stable on reload. + created_at = time.time() + + def _persist() -> list[dict]: + records = [] + for image in result["images"]: + records.append(image_gallery.save(image, { + "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": created_at, + })) + return records + try: - record, image_b64 = await asyncio.to_thread(image_gallery.save, result["image"], meta) + records = await asyncio.to_thread(_persist) 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), - ) + return DiffusionGenerateResponse(images = [GalleryImage(**r) for r in records]) @studio_router.get("/images/gallery", response_model = GalleryListResponse) diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index e72ec18d5f..d6085a4b43 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -122,7 +122,8 @@ class _FakePipe: def __call__(self, **kwargs): self.last_kwargs = kwargs - return types.SimpleNamespace(images = [_FakeImage()]) + n = kwargs.get("num_images_per_prompt", 1) + return types.SimpleNamespace(images = [_FakeImage() for _ in range(n)]) class _FakePipeline: @@ -193,11 +194,15 @@ def test_load_generate_unload_gguf(fake_runtime, tmp_path): gen = backend.generate(prompt = "a sloth", width = 512, height = 512, steps = 4, guidance = 3.0) assert gen["seed"] == 4242 # random seed reported back 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 + assert len(gen["images"]) == 1 # PIL images handed to the route for persistence gen2 = backend.generate(prompt = "again", seed = 99) assert gen2["seed"] == 99 + # batch_size produces that many images in one call, all sharing the seed. + batch = backend.generate(prompt = "batch", seed = 7, batch_size = 3) + assert len(batch["images"]) == 3 and batch["seed"] == 7 + assert backend.unload()["loaded"] is False assert backend.is_loaded is False diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index dda2edb96f..0e941f5a77 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -51,12 +51,16 @@ class _FakeBackend: "error": None, } - def generate(self, *, seed = None, **kwargs): + def generate(self, *, seed = None, batch_size = 1, **kwargs): if not self.loaded: raise RuntimeError("No diffusion model is loaded.") - # 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"} + # The real backend returns the PIL images; the route persists them. The + # fake returns sentinels since image_gallery is stubbed in the fixture. + return { + "images": [object() for _ in range(batch_size)], + "seed": seed if seed is not None else 4242, + "repo_id": "x/z-image", + } def unload(self): self.loaded = False @@ -97,7 +101,7 @@ def client(monkeypatch, tmp_path): (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) + return record def _clear(): n = len(store) @@ -137,10 +141,10 @@ 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["image_b64"] - # The persisted record carries the full recipe back. - img = gbody["image"] + # One persisted record carrying the full recipe back. + images = gen.json()["images"] + assert len(images) == 1 + img = images[0] assert img["seed"] == 7 and img["prompt"] == "a sloth" and img["id"] # The image is now listable, fetchable, and deletable. @@ -155,6 +159,20 @@ def test_load_generate_status_unload_roundtrip(client): assert client.get("/api/inference/images/status").json()["loaded"] is False +def test_generate_batch_size_persists_each_image(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", "batch_size": 3, "seed": 5}, + ) + assert resp.status_code == 200 + images = resp.json()["images"] + assert len(images) == 3 + assert all(i["seed"] == 5 for i in images) # the batch shares one seed + assert len({i["id"] for i in images}) == 3 # but each is a distinct record + assert len(client.get("/api/inference/images/gallery").json()["images"]) == 3 + + 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 8228aedc1d..52e721dad6 100644 --- a/studio/backend/tests/test_image_gallery.py +++ b/studio/backend/tests/test_image_gallery.py @@ -44,7 +44,7 @@ def _meta(**over): def test_save_embeds_recipe_and_round_trips(): - record, _b64 = gallery.save(_img(), _meta()) + record = 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. @@ -60,20 +60,20 @@ def test_save_embeds_recipe_and_round_trips(): 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 = 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")) + 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")) + 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 diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index 960f6379c4..aad4b0cc78 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -39,6 +39,7 @@ export interface DiffusionGenerateRequest { steps?: number; guidance?: number; seed?: number; + batch_size?: number; } // A persisted image's full generation recipe (also embedded in the PNG). @@ -57,9 +58,7 @@ export interface GalleryImage { } export interface DiffusionGenerateResponse { - image_b64: string; - mime: string; - image: GalleryImage; + images: GalleryImage[]; } async function parseJson(response: Response): Promise { diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 814fbf85d0..97e007a494 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -321,8 +321,13 @@ export function ImagesPage() { const [steps, setSteps] = useState(9); const [guidance, setGuidance] = useState(0.0); const [seed, setSeed] = useState(""); + // Batch size = images per forward pass (VRAM-heavy); count = sequential loops. + const [batchSize, setBatchSize] = useState(1); + const [count, setCount] = useState(1); const [busy, setBusy] = useState(null); + // {done, total} while a multi-run generation is in flight (for the button). + const [genProgress, setGenProgress] = useState<{ done: number; total: number } | null>(null); const [status, setStatus] = useState(null); // Records come from the backend (durable); srcById maps each id to its object // URL (loaded images) or data URL (the one just generated). @@ -515,40 +520,48 @@ export function ImagesPage() { toast.error("Prompt is empty"); return; } - let parsedSeed: number | undefined; + // Resolve a base seed up front. With an explicit seed the run is fully + // reproducible; with a random one we still pick a concrete base now so each + // sequential image gets a distinct, reproducible seed (base + i). + let baseSeed: number; if (seed.trim()) { const n = Number(seed); - if (!Number.isInteger(n) || n < 0) { + if (!Number.isInteger(n) || n < 0 || n > Number.MAX_SAFE_INTEGER) { toast.error("Seed must be a non-negative integer"); return; } - parsedSeed = n; + baseSeed = n; + } else { + baseSeed = Math.floor(Math.random() * 2 ** 32); } setBusy("generating"); + setGenProgress({ done: 0, total: count }); try { - const res = await generateDiffusionImage({ - prompt: prompt.trim(), - negative_prompt: negativePrompt.trim() || undefined, - width: resolution.w, - height: resolution.h, - steps, - guidance, - seed: parsedSeed, - }); - // 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); + for (let i = 0; i < count; i++) { + const res = await generateDiffusionImage({ + prompt: prompt.trim(), + negative_prompt: negativePrompt.trim() || undefined, + width: resolution.w, + height: resolution.h, + steps, + guidance, + seed: baseSeed + i, + batch_size: batchSize, + }); + // Prepend this run's records (newest first) and load their blobs. + setImages((prev) => [...res.images, ...prev]); + if (res.images[0]) setSelectedId(res.images[0].id); + res.images.forEach((image) => void ensureSrc(image)); + setGenProgress({ done: i + 1, total: count }); + } } catch (err) { toast.error(err instanceof Error ? err.message : "Image generation failed"); } finally { setBusy(null); + setGenProgress(null); } - }, [prompt, negativePrompt, resolution, steps, guidance, seed]); + }, [prompt, negativePrompt, resolution, steps, guidance, seed, batchSize, count, ensureSrc]); return (
@@ -635,9 +648,30 @@ export function ImagesPage() { /> + + +