ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add batch size and sequential count to image generation with reproducible per-run seeds

Browse source
This commit is contained in:
oobabooga 2026-06-24 18:02:38 -03:00
commit eced3620fa
9 changed files with 139 additions and 75 deletions
Download patch file Download diff file Expand all files Collapse all files

12
studio/backend/core/inference/diffusion.py
View file

@ -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:

12
studio/backend/core/inference/image_gallery.py
View file

@ -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]:

11
studio/backend/models/inference.py
View file

@ -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):

43
studio/backend/routes/inference.py
View file

@ -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)

9
studio/backend/tests/test_diffusion_backend.py
View file

@ -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

36
studio/backend/tests/test_diffusion_routes.py
View file

@ -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

10
studio/backend/tests/test_image_gallery.py
View file

@ -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

5
studio/frontend/src/features/images/api.ts
View file

@ -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<T>(response: Response): Promise<T> {

76
studio/frontend/src/features/images/images-page.tsx
View file

@ -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<Busy>(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<DiffusionStatus | null>(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 (
<div className="flex h-full min-h-0 min-w-0 flex-1 flex-col overflow-hidden">
@ -635,9 +648,30 @@ export function ImagesPage() {
/>
</Field>
<SliderField
label="Batch size"
hint="Images generated together in one forward pass (VRAM-heavy). They share the run's seed but each comes out different."
value={batchSize}
min={1}
max={32}
step={1}
onChange={setBatchSize}
/>
<SliderField
label="Sequential count"
hint="Repeat generation this many times in a loop. Each run advances the seed (base + i), so images differ and stay reproducible."
value={count}
min={1}
max={128}
step={1}
onChange={setCount}
/>
<Button onClick={handleGenerate} disabled={busy !== null || !status?.loaded}>
{busy === "generating" ? <Spinner className="mr-2 size-4" /> : null}
Generate
{busy === "generating" && genProgress && genProgress.total > 1
? `Generating ${genProgress.done}/${genProgress.total}…`
: "Generate"}
</Button>
</SectionCard>