From 571d506c206dcb8bbeac5c86e88e7cdc848ae2c8 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Wed, 24 Jun 2026 18:47:33 -0300 Subject: [PATCH] Refine image generation UI: aspect ratio sliders, info hints, action toolbar, and timestamped export filenames --- studio/backend/models/inference.py | 1 + studio/backend/routes/inference.py | 5 +- studio/frontend/src/features/images/api.ts | 1 + .../src/features/images/images-page.tsx | 227 ++++++++++++------ 4 files changed, 164 insertions(+), 70 deletions(-) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 0a5e18758a..e8b800cbf8 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1740,6 +1740,7 @@ class GalleryImage(BaseModel): steps: int = Field(..., description = "Denoising steps") guidance: float = Field(..., description = "Guidance scale") seed: int = Field(..., description = "Seed used") + batch_index: int = Field(0, description = "Position within its batch (0-based)") model: Optional[str] = Field(None, description = "Model repo id that produced it") created_at: float = Field(..., description = "Creation time (epoch seconds)") diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 96c33f8528..6f7772417a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10117,7 +10117,7 @@ async def generate_diffusion_image( def _persist() -> list[dict]: records = [] - for image in result["images"]: + for index, image in enumerate(result["images"]): records.append(image_gallery.save(image, { "prompt": request.prompt, "negative_prompt": request.negative_prompt, @@ -10126,6 +10126,9 @@ async def generate_diffusion_image( "steps": request.steps, "guidance": request.guidance, "seed": result["seed"], + # Position within the batch: images here share a seed + timestamp, + # so the export filename needs this to stay unique. + "batch_index": index, "model": result.get("repo_id"), "created_at": created_at, })) diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index aad4b0cc78..4ab76b0360 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -53,6 +53,7 @@ export interface GalleryImage { steps: number; guidance: number; seed: number; + batch_index: number; model: string | null; created_at: number; } diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 97e007a494..64d4661ac0 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -3,6 +3,7 @@ import { useCallback, useEffect, useMemo, useRef, useState, type ReactNode } from "react"; import { + ArrowLeftRightIcon, ArrowReloadHorizontalIcon, Delete02Icon, Download01Icon, @@ -68,21 +69,35 @@ const MODELS: ModelOption[] = [ { id: MODEL.repo_id, name: MODEL.label, description: "Text-to-image · GGUF", isGguf: true }, ]; -// 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: "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 }, -]; +// Common aspect ratios (landscape-oriented; the Flip button gives the portrait +// mirror). Picking one locks the W:H proportion; the sliders size it, and the +// generation snaps to Z-Image's ~1-megapixel-trained 16px grid. +const ASPECT_RATIOS: Record = { + "1:1": [1, 1], + "3:2": [3, 2], + "4:3": [4, 3], + "16:9": [16, 9], + "21:9": [21, 9], +}; +const ASPECT_OPTIONS = ["custom", ...Object.keys(ASPECT_RATIOS)]; + +// Z-Image accepts 256–2048, in multiples of 16. Snap any value into range. +const MIN_DIM = 256; +const MAX_DIM = 2048; +function snapDim(value: number): number { + if (!Number.isFinite(value)) return 1024; + return Math.min(MAX_DIM, Math.max(MIN_DIM, Math.round(value / 16) * 16)); +} + +// The ratio key (compared by long:short, so it survives orientation) matching +// width/height, plus whether the current orientation is portrait. +function matchAspect(width: number, height: number): { key: string; portrait: boolean } { + const target = Math.max(width, height) / Math.min(width, height); + const found = Object.entries(ASPECT_RATIOS).find( + ([, [a, b]]) => Math.abs(target - a / b) < 0.01, + ); + return { key: found ? found[0] : "custom", portrait: height > width }; +} // 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 @@ -98,10 +113,24 @@ const galleryCache: { inflight: Set; } = { images: [], selectedId: null, quant: null, srcById: new Map(), inflight: new Set() }; -function downloadImage(src: string, seed: number) { +// Export filename: app name, a compact sortable timestamp, and the seed. A batch +// shares the seed + timestamp, so a "_" suffix is added past the first image. +// e.g. Unsloth_20260624-143005_123.png / Unsloth_20260624-143005_123_1.png +// (the full recipe is embedded in the PNG regardless). +function exportFilename(image: GalleryImage): string { + const d = new Date(image.created_at * 1000); + const p = (n: number) => String(n).padStart(2, "0"); + const stamp = + `${d.getFullYear()}${p(d.getMonth() + 1)}${p(d.getDate())}` + + `-${p(d.getHours())}${p(d.getMinutes())}${p(d.getSeconds())}`; + const suffix = image.batch_index > 0 ? `_${image.batch_index}` : ""; + return `Unsloth_${stamp}_${image.seed}${suffix}.png`; +} + +function downloadImage(src: string, image: GalleryImage) { const link = document.createElement("a"); link.href = src; - link.download = `unsloth-${seed}.png`; + link.download = exportFilename(image); link.click(); } @@ -275,7 +304,7 @@ function RecipePopover({ return ( - @@ -315,7 +344,12 @@ export function ImagesPage() { "a tiny ginger sloth coding in a sunlit treehouse, photorealistic", ); const [negativePrompt, setNegativePrompt] = useState(""); - const [resolutionIdx, setResolutionIdx] = useState(0); + // width/height are the source of truth; `aspect` locks their proportion + // ("custom" = free) and `portrait` tracks orientation, so Flip keeps the lock. + const [width, setWidth] = useState(1024); + const [height, setHeight] = useState(1024); + const [aspect, setAspect] = useState("1:1"); + const [portrait, setPortrait] = useState(false); // 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); @@ -352,7 +386,6 @@ export function ImagesPage() { galleryCache.quant = quant; }, [images, selectedId, quant]); - const resolution = RESOLUTIONS[resolutionIdx]; const selected = useMemo( () => images.find((i) => i.id === selectedId) ?? images[0] ?? null, [images, selectedId], @@ -415,11 +448,44 @@ export function ImagesPage() { 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); + setWidth(image.width); + setHeight(image.height); + const m = matchAspect(image.width, image.height); + setAspect(m.key); + setPortrait(m.portrait); toast.success("Settings restored to inputs"); }, []); + // Size controls: a locked aspect ratio keeps the paired dimension in step as + // you drag a slider; "custom" frees both. Flip swaps width and height and + // keeps the lock (the ratio just applies in the other orientation). + // h/w multiplier for the locked ratio [a,b] (a:b = long:short): landscape puts + // the long side on width (h = w*b/a), portrait puts it on height (h = w*a/b). + const ratioHW = (a: number, b: number) => (portrait ? a / b : b / a); + const changeAspect = (key: string) => { + setAspect(key); + if (key === "custom") return; + const [a, b] = ASPECT_RATIOS[key]; + setHeight(snapDim(width * ratioHW(a, b))); + }; + const changeWidth = (v: number) => { + setWidth(v); + if (aspect === "custom") return; + const [a, b] = ASPECT_RATIOS[aspect]; + setHeight(snapDim(v * ratioHW(a, b))); + }; + const changeHeight = (v: number) => { + setHeight(v); + if (aspect === "custom") return; + const [a, b] = ASPECT_RATIOS[aspect]; + setWidth(snapDim(v / ratioHW(a, b))); + }; + const flipDimensions = () => { + setWidth(height); + setHeight(width); + setPortrait((p) => !p); + }; + const refreshStatus = useCallback(async () => { try { setStatus(await getDiffusionStatus()); @@ -535,6 +601,10 @@ export function ImagesPage() { baseSeed = Math.floor(Math.random() * 2 ** 32); } + // Snap custom dims to the model's grid so a half-typed value can't 422. + const w = snapDim(width); + const h = snapDim(height); + setBusy("generating"); setGenProgress({ done: 0, total: count }); try { @@ -542,8 +612,8 @@ export function ImagesPage() { const res = await generateDiffusionImage({ prompt: prompt.trim(), negative_prompt: negativePrompt.trim() || undefined, - width: resolution.w, - height: resolution.h, + width: w, + height: h, steps, guidance, seed: baseSeed + i, @@ -561,7 +631,7 @@ export function ImagesPage() { setBusy(null); setGenProgress(null); } - }, [prompt, negativePrompt, resolution, steps, guidance, seed, batchSize, count, ensureSrc]); + }, [prompt, negativePrompt, width, height, steps, guidance, seed, batchSize, count, ensureSrc]); return (
@@ -594,7 +664,7 @@ export function ImagesPage() {