RomanNum3ral
/
ai-toolkit
mirror of https://github.com/ostris/ai-toolkit.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Switch to thumbnails and thumbnail creation on the sample grid page until clicked.

This commit is contained in:
Jaret Burkett 2026-07-25 08:40:33 -06:00
parent e00f3791e2
commit efb58c8641
7 changed files with 105 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

52
toolkit/config_modules.py
View File

@ -1227,6 +1227,58 @@ class GenerateImageConfig:
# join with folder
return os.path.join(self.output_folder, filename)
def save_image_atomic(self, image, count: int = 0, max_count=0):
# write into a hidden tmp subfolder, then atomically move into place so
# watchers (UI/CDN) never see and cache a partially written file. Wraps
# self.save_image so it also covers models that replace that function.
real_folder = self.output_folder
tmp_folder = os.path.join(real_folder, '.tmp')
os.makedirs(tmp_folder, exist_ok=True)
self.output_folder = tmp_folder
try:
self.save_image(image, count, max_count)
finally:
self.output_folder = real_folder
files = os.listdir(tmp_folder)
# thumbs move into place first so they already exist when the media
# file appears in the samples folder
thumbs_folder = os.path.join(real_folder, '.thumbs')
for file in files:
tmp_thumb = os.path.join(tmp_folder, file + '.thumb')
try:
if self._generate_thumbnail(os.path.join(tmp_folder, file), tmp_thumb):
os.makedirs(thumbs_folder, exist_ok=True)
os.replace(tmp_thumb, os.path.join(thumbs_folder, file + '.jpg'))
except Exception as e:
print(f"Failed to generate thumbnail for {file}: {e}")
for file in files:
os.replace(os.path.join(tmp_folder, file), os.path.join(real_folder, file))
def _generate_thumbnail(self, media_path, thumb_path):
# 300x300 center-cropped 90% jpg. Returns True if one was written.
from PIL import Image as PILImage
ext = os.path.splitext(media_path)[1].lower()
img = None
if ext in ['.png', '.jpg', '.jpeg', '.webp', '.gif', '.bmp']:
img = PILImage.open(media_path) # animated formats open on the first frame
elif ext == '.mp4':
import cv2
cap = cv2.VideoCapture(media_path)
ok, frame = cap.read()
cap.release()
if ok:
img = PILImage.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
if img is None:
return False
img = img.convert('RGB')
w, h = img.size
side = min(w, h)
left = (w - side) // 2
top = (h - side) // 2
img = img.crop((left, top, left + side, top + side)).resize((300, 300), PILImage.LANCZOS)
img.save(thumb_path, format='JPEG', quality=90)
return True
def save_image(self, image, count: int = 0, max_count=0):
# make parent dirs
os.makedirs(self.output_folder, exist_ok=True)

2
toolkit/models/base_model.py
View File

@ -666,7 +666,7 @@ class BaseModel:
extra,
)
gen_config.save_image(img, i)
gen_config.save_image_atomic(img, i)
gen_config.log_image(img, i)
self._after_sample_image(i, len(image_configs))
flush()

2
toolkit/stable_diffusion_model.py
View File

@ -1715,7 +1715,7 @@ class StableDiffusion:
generator=generator,
).images[0]
gen_config.save_image(img, i)
gen_config.save_image_atomic(img, i)
gen_config.log_image(img, i)
self._after_sample_image(i, len(image_configs))
flush()

13
ui/cron/fileServer.ts
View File

@ -118,7 +118,7 @@ async function serveFile(req: http.IncomingMessage, res: http.ServerResponse, pr
.map(decodeURIComponent)
.join('/');
const resolvedFilePath = path.resolve(decodedFilePath);
let resolvedFilePath = path.resolve(decodedFilePath);
const roots = await getRoots();
const allowedDirs = isImg ? [roots.datasets, roots.training, roots.data] : [roots.datasets, roots.training];
const isAllowed = allowedDirs.some(
@ -131,6 +131,17 @@ async function serveFile(req: http.IncomingMessage, res: http.ServerResponse, pr
return;
}
// ?thumb=1 serves the pre-generated 300x300 jpg from the sibling .thumbs
// folder (<name>.<ext>.jpg) when it exists; otherwise falls through to
// the full file exactly as before. Mirrors the Next.js /api/img route.
if (isImg && new URL(req.url || '', 'http://localhost').searchParams.has('thumb')) {
const thumbPath = path.join(path.dirname(resolvedFilePath), '.thumbs', path.basename(resolvedFilePath) + '.jpg');
const thumbStat = await fs.promises.stat(thumbPath).catch(() => null);
if (thumbStat && thumbStat.isFile()) {
resolvedFilePath = thumbPath;
}
}
let stat: fs.Stats;
try {
stat = await fs.promises.stat(resolvedFilePath);

13
ui/src/app/api/img/[...imagePath]/route.ts
View File

@ -45,7 +45,7 @@ export async function GET(request: NextRequest, { params }: { params: { imagePat
// Security check: resolve the path so any `..` segments are collapsed,
// then ensure it's still under an allowed root. (Plain `.includes('..')`
// false-positives on filenames that contain `..` as text, e.g. an ellipsis.)
const resolved = path.resolve(filepath);
let resolved = path.resolve(filepath);
const isAllowed = allowedDirs.some(
allowedDir => resolved === allowedDir || resolved.startsWith(allowedDir + path.sep),
);
@ -55,6 +55,17 @@ export async function GET(request: NextRequest, { params }: { params: { imagePat
return new NextResponse('Access denied', { status: 403 });
}
// ?thumb=1 serves the pre-generated 300x300 jpg from the sibling .thumbs
// folder (<name>.<ext>.jpg) when it exists; otherwise falls through to
// the full file exactly as before.
if (request.nextUrl.searchParams.has('thumb')) {
const thumbPath = path.join(path.dirname(resolved), '.thumbs', path.basename(resolved) + '.jpg');
const thumbStat = await fs.promises.stat(thumbPath).catch(() => null);
if (thumbStat && thumbStat.isFile()) {
resolved = thumbPath;
}
}
// Bail out early if the client already gave up
if (request.signal.aborted) {
return new NextResponse(null, { status: 499 });

8
ui/src/app/api/jobs/[jobID]/samples/route.ts
View File

@ -25,8 +25,12 @@ export async function GET(request: NextRequest, { params }: { params: { jobID: s
return NextResponse.json({ samples: [] });
}
// find all img (png, jpg, jpeg) files in the samples folder
const samples = (await fs.promises.readdir(samplesFolder))
// find all img (png, jpg, jpeg) files in the samples folder. Thumbnails
// live in the hidden .thumbs subfolder (and partial writes in .tmp) — the
// isFile check keeps those directories out even if their names ever match.
const samples = (await fs.promises.readdir(samplesFolder, { withFileTypes: true }))
.filter(entry => entry.isFile())
.map(entry => entry.name)
.filter(file => {
return file.endsWith('.png') || file.endsWith('.jpg') || file.endsWith('.jpeg') || file.endsWith('.webp') || file.endsWith('.mp4') || file.endsWith('mp3') || file.endsWith('wav') || file.endsWith('flac') || file.endsWith('ogg');
})

28
ui/src/components/SampleImageCard.tsx
View File

@ -32,10 +32,11 @@ const SampleImageCard: React.FC<SampleImageCardProps> = ({
const [isVisible, setIsVisible] = useState(false);
const [loaded, setLoaded] = useState(false);
const [blobUrl, setBlobUrl] = useState<string | null>(null);
// videos with no pre-generated thumb (older samples) fall back to the <video> element
const [videoFallback, setVideoFallback] = useState(false);
const isItAudio = isAudio(imageUrl);
const isItVideo = isVideo(imageUrl);
const isImageType = !isItAudio && !isItVideo;
// Observe both enter and exit
useEffect(() => {
@ -66,7 +67,7 @@ const SampleImageCard: React.FC<SampleImageCardProps> = ({
// the element unmounts). A short debounce skips requests entirely during fast
// scrolls where the card is only briefly visible.
useEffect(() => {
if (!isImageType) return;
if (isItAudio) return;
if (!isVisible) return;
const controller = new AbortController();
@ -74,13 +75,25 @@ const SampleImageCard: React.FC<SampleImageCardProps> = ({
let objectUrl: string | null = null;
const timer = window.setTimeout(() => {
fetch(`/api/img/${encodeURIComponent(imageUrl)}`, { signal: controller.signal })
// ?thumb=1: the server sends the small pre-generated thumbnail when one
// exists, otherwise the full file. Videos without a thumb come back as
// video/* — abort the transfer and render the <video> element instead.
fetch(`/api/img/${encodeURIComponent(imageUrl)}?thumb=1`, { signal: controller.signal })
.then(r => {
if (!r.ok) throw new Error(`HTTP ${r.status}`);
const contentType = r.headers.get('content-type') || '';
if (isItVideo && !contentType.startsWith('image/')) {
controller.abort();
if (!cancelled) {
setVideoFallback(true);
setLoaded(true);
}
return null;
}
return r.blob();
})
.then(blob => {
if (cancelled) return;
if (cancelled || !blob) return;
objectUrl = URL.createObjectURL(blob);
setBlobUrl(objectUrl);
setLoaded(true);
@ -97,15 +110,16 @@ const SampleImageCard: React.FC<SampleImageCardProps> = ({
if (objectUrl) URL.revokeObjectURL(objectUrl);
setBlobUrl(null);
setLoaded(false);
setVideoFallback(false);
};
}, [isVisible, isImageType, imageUrl]);
}, [isVisible, isItAudio, isItVideo, imageUrl]);
return (
<div className={`flex flex-col ${className}`}>
<div ref={cardRef} className="relative w-full cursor-pointer" style={{ paddingBottom: '100%' }} onClick={onClick}>
<div
className={`absolute inset-0 rounded-t-lg shadow-md bg-gray-900 ${
isVisible && isImageType && !loaded ? 'animate-pulse' : ''
isVisible && !isItAudio && !loaded ? 'animate-pulse' : ''
}`}
>
{isVisible ? (
@ -120,7 +134,7 @@ const SampleImageCard: React.FC<SampleImageCardProps> = ({
}}
/>
</div>
) : isItVideo ? (
) : isItVideo && videoFallback ? (
<video
ref={videoRef}
src={`/api/img/${encodeURIComponent(imageUrl)}`}