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

Drastically improve the loading speed of images in the ui by using a custom loader and abort controller to abort when images leave the view.

This commit is contained in:
Jaret Burkett 2026-05-25 10:04:25 -06:00
parent 7354def271
commit 083236a2a7
5 changed files with 203 additions and 57 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -2,6 +2,7 @@
import { NextRequest, NextResponse } from 'next/server';
import fs from 'fs';
import path from 'path';
import { Readable } from 'stream';
import { getDatasetsRoot, getTrainingFolder, getDataRoot } from '@/server/settings';
const contentTypeMap: { [key: string]: string } = {
@ -54,6 +55,11 @@ export async function GET(request: NextRequest, { params }: { params: { imagePat
return new NextResponse('Access denied', { status: 403 });
}
// Bail out early if the client already gave up
if (request.signal.aborted) {
return new NextResponse(null, { status: 499 });
}
// Stat file (async)
const stat = await fs.promises.stat(resolved).catch(() => null);
if (!stat || !stat.isFile()) {
@ -63,6 +69,40 @@ export async function GET(request: NextRequest, { params }: { params: { imagePat
const ext = path.extname(resolved).toLowerCase();
const contentType = contentTypeMap[ext] || 'application/octet-stream';
// Weak ETag from inode/size/mtime — cheap and stable enough for revalidation
const etag = `W/"${stat.ino.toString(36)}-${stat.size.toString(36)}-${stat.mtimeMs.toString(36)}"`;
const cacheControl = 'public, max-age=86400, immutable';
const ifNoneMatch = request.headers.get('if-none-match');
if (ifNoneMatch && ifNoneMatch === etag) {
return new NextResponse(null, {
status: 304,
headers: {
ETag: etag,
'Cache-Control': cacheControl,
},
});
}
const buildBody = (start?: number, end?: number) => {
const nodeStream =
start !== undefined && end !== undefined
? fs.createReadStream(resolved, { start, end })
: fs.createReadStream(resolved);
// Wire client disconnect → destroy the file stream so we don't keep
// reading bytes for a request the browser has already cancelled.
const onAbort = () => nodeStream.destroy();
if (request.signal.aborted) {
nodeStream.destroy();
} else {
request.signal.addEventListener('abort', onAbort, { once: true });
}
nodeStream.once('close', () => request.signal.removeEventListener('abort', onAbort));
return Readable.toWeb(nodeStream) as unknown as ReadableStream;
};
// Support range requests for video/audio seeking
const rangeHeader = request.headers.get('range');
if (rangeHeader) {
@ -71,49 +111,26 @@ export async function GET(request: NextRequest, { params }: { params: { imagePat
const end = parts[1] ? parseInt(parts[1], 10) : stat.size - 1;
const chunkSize = end - start + 1;
const stream = fs.createReadStream(resolved, { start, end });
const readable = new ReadableStream({
start(controller) {
stream.on('data', chunk => controller.enqueue(chunk));
stream.on('end', () => controller.close());
stream.on('error', err => controller.error(err));
},
cancel() {
stream.destroy();
},
});
return new NextResponse(readable as any, {
return new NextResponse(buildBody(start, end) as any, {
status: 206,
headers: {
'Content-Range': `bytes ${start}-${end}/${stat.size}`,
'Accept-Ranges': 'bytes',
'Content-Length': String(chunkSize),
'Content-Type': contentType,
'Cache-Control': 'public, max-age=86400',
'Cache-Control': cacheControl,
ETag: etag,
},
});
}
// Stream the file instead of buffering it entirely
const stream = fs.createReadStream(resolved);
const readable = new ReadableStream({
start(controller) {
stream.on('data', chunk => controller.enqueue(chunk));
stream.on('end', () => controller.close());
stream.on('error', err => controller.error(err));
},
cancel() {
stream.destroy();
},
});
return new NextResponse(readable as any, {
return new NextResponse(buildBody() as any, {
headers: {
'Content-Type': contentType,
'Content-Length': String(stat.size),
'Cache-Control': 'public, max-age=86400',
'Cache-Control': cacheControl,
'Accept-Ranges': 'bytes',
ETag: etag,
},
});
} catch (error) {

1
ui/src/app/datasets/[datasetName]/page.tsx
View File

@ -150,6 +150,7 @@ export default function DatasetPage({ params }: { params: { datasetName: string
onDelete={() => refreshImageList(datasetName)}
onImageClick={() => setSelectedImgPath(img.img_path)}
captionRefreshKey={captionRefreshKeys[img.img_path] || 0}
observerRoot={scrollParent}
/>
);
}}

95
ui/src/components/DatasetImageCard.tsx
View File

@ -16,6 +16,8 @@ interface DatasetImageCardProps {
onDelete?: () => void;
onImageClick?: () => void;
captionRefreshKey?: number;
observerRoot?: Element | null;
rootMargin?: string;
}
const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
@ -27,13 +29,88 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
onDelete = () => {},
onImageClick,
captionRefreshKey = 0,
observerRoot = null,
rootMargin = '200px 0px',
}) => {
const [loaded, setLoaded] = useState<boolean>(false);
const [showAudioPlayer, setShowAudioPlayer] = useState(true);
const [pollTick, setPollTick] = useState(0);
const [blobUrl, setBlobUrl] = useState<string | null>(null);
const [isVisible, setIsVisible] = useState(false);
const cardRef = useRef<HTMLDivElement>(null);
const isItAVideo = isVideo(imageUrl);
const isItAudio = isAudio(imageUrl);
const isItImage = !isItAVideo && !isItAudio;
// Track actual viewport visibility — Virtuoso keeps a buffer of cards mounted
// outside the visible region, so we can't rely on mount/unmount alone.
useEffect(() => {
const el = cardRef.current;
if (!el) return;
const observer = new IntersectionObserver(
entries => {
for (const entry of entries) {
if (entry.target === el) {
setIsVisible(entry.isIntersecting);
}
}
},
{
root: observerRoot ?? null,
threshold: 0.01,
rootMargin,
},
);
observer.observe(el);
return () => observer.disconnect();
}, [observerRoot, rootMargin]);
// Drive image loads through fetch + AbortController so scrolling past actually
// cancels in-flight requests. Debounced 80ms so fast scroll-throughs never
// start a request.
useEffect(() => {
if (!isItImage) return;
if (!isVisible) return;
const controller = new AbortController();
let cancelled = false;
let objectUrl: string | null = null;
const timer = window.setTimeout(() => {
fetch(`/api/img/${encodeURIComponent(imageUrl)}`, { signal: controller.signal })
.then(r => {
if (!r.ok) throw new Error(`HTTP ${r.status}`);
return r.blob();
})
.then(blob => {
if (cancelled) return;
objectUrl = URL.createObjectURL(blob);
setBlobUrl(objectUrl);
setLoaded(true);
})
.catch(err => {
if (err?.name !== 'AbortError') console.error('Dataset image fetch failed:', err);
});
}, 80);
return () => {
cancelled = true;
clearTimeout(timer);
controller.abort();
if (objectUrl) URL.revokeObjectURL(objectUrl);
setBlobUrl(null);
setLoaded(false);
};
}, [imageUrl, isItImage, isVisible]);
const combinedRefreshKey = captionRefreshKey + pollTick;
const { caption: fetchedCaption, isLoaded: isCaptionLoaded } = useCaptionBatch(imageUrl, combinedRefreshKey);
const { caption: fetchedCaption, isLoaded: isCaptionLoaded } = useCaptionBatch(
isVisible ? imageUrl : null,
combinedRefreshKey,
);
const [caption, setCaption] = useState<string>('');
const [savedCaption, setSavedCaption] = useState<string>('');
@ -90,10 +167,6 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
};
}, []);
const handleLoad = (): void => {
setLoaded(true);
};
const handleKeyDown = (e: KeyboardEvent<HTMLTextAreaElement>): void => {
if (e.key === 'Enter' && !e.shiftKey) {
e.preventDefault();
@ -108,12 +181,8 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
const isCaptionCurrent = caption.trim() === savedCaption;
const isItAVideo = isVideo(imageUrl);
const isItAudio = isAudio(imageUrl);
const isItImage = !isItAVideo && !isItAudio;
return (
<div className={`flex flex-col ${className}`}>
<div ref={cardRef} className={`flex flex-col ${className}`}>
<div className="relative w-full" style={{ paddingBottom: '100%' }}>
<div
className={classNames('absolute inset-0 rounded-t-lg shadow-md bg-gray-900', {
@ -148,11 +217,11 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
{isItAudio && showAudioPlayer && (
<AudioPlayer src={`/api/img/${encodeURIComponent(imageUrl)}`} title={imageUrl.replace(/^.*[\\/]/, '')} />
)}
{isItImage && (
{isItImage && blobUrl && (
<img
src={`/api/img/${encodeURIComponent(imageUrl)}`}
src={blobUrl}
alt={alt}
onLoad={handleLoad}
decoding="async"
onClick={onImageClick}
className={classNames('w-full h-full object-contain transition-opacity duration-300', {
'opacity-100': loaded,

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

@ -31,6 +31,11 @@ const SampleImageCard: React.FC<SampleImageCardProps> = ({
const videoRef = useRef<HTMLVideoElement | null>(null);
const [isVisible, setIsVisible] = useState(false);
const [loaded, setLoaded] = useState(false);
const [blobUrl, setBlobUrl] = useState<string | null>(null);
const isItAudio = isAudio(imageUrl);
const isItVideo = isVideo(imageUrl);
const isImageType = !isItAudio && !isItVideo;
// Observe both enter and exit
useEffect(() => {
@ -56,9 +61,44 @@ const SampleImageCard: React.FC<SampleImageCardProps> = ({
return () => observer.disconnect();
}, [observerRoot, rootMargin]);
const handleLoad = () => setLoaded(true);
// Drive image loads through fetch + AbortController so scrolling past actually
// cancels in-flight requests (browsers don't reliably cancel <img> fetches when
// the element unmounts). A short debounce skips requests entirely during fast
// scrolls where the card is only briefly visible.
useEffect(() => {
if (!isImageType) return;
if (!isVisible) return;
const isImageType = !isAudio(imageUrl) && !isVideo(imageUrl);
const controller = new AbortController();
let cancelled = false;
let objectUrl: string | null = null;
const timer = window.setTimeout(() => {
fetch(`/api/img/${encodeURIComponent(imageUrl)}`, { signal: controller.signal })
.then(r => {
if (!r.ok) throw new Error(`HTTP ${r.status}`);
return r.blob();
})
.then(blob => {
if (cancelled) return;
objectUrl = URL.createObjectURL(blob);
setBlobUrl(objectUrl);
setLoaded(true);
})
.catch(err => {
if (err?.name !== 'AbortError') console.error('Sample image fetch failed:', err);
});
}, 80);
return () => {
cancelled = true;
clearTimeout(timer);
controller.abort();
if (objectUrl) URL.revokeObjectURL(objectUrl);
setBlobUrl(null);
setLoaded(false);
};
}, [isVisible, isImageType, imageUrl]);
return (
<div className={`flex flex-col ${className}`}>
@ -69,7 +109,7 @@ const SampleImageCard: React.FC<SampleImageCardProps> = ({
}`}
>
{isVisible ? (
isAudio(imageUrl) ? (
isItAudio ? (
<div className="w-full h-full flex items-center justify-center bg-gray-900">
<img
src={`/api/audio/art/${encodeURIComponent(imageUrl)}`}
@ -80,31 +120,28 @@ const SampleImageCard: React.FC<SampleImageCardProps> = ({
}}
/>
</div>
) : isVideo(imageUrl) ? (
) : isItVideo ? (
<video
ref={videoRef}
src={`/api/img/${encodeURIComponent(imageUrl)}`}
className="w-full h-full object-cover"
preload="none"
onLoad={handleLoad}
playsInline
muted
loop
autoPlay
controls={false}
/>
) : (
) : blobUrl ? (
<img
src={`/api/img/${encodeURIComponent(imageUrl)}`}
src={blobUrl}
alt={alt}
onLoad={handleLoad}
loading="lazy"
decoding="async"
className={`w-full h-full object-cover transition-opacity duration-300 ${
loaded ? 'opacity-100' : 'opacity-0'
}`}
/>
)
) : null
) : null}
{children && isVisible && <div className="absolute inset-0 flex items-center justify-center">{children}</div>}

32
ui/src/hooks/useCaptionBatch.tsx
View File

@ -46,13 +46,33 @@ async function flush() {
if (pending.size > 0) scheduleFlush();
}
function requestCaption(path: string): Promise<string> {
function requestCaption(path: string, signal?: AbortSignal): Promise<string> {
return new Promise((resolve, reject) => {
if (signal?.aborted) {
reject(new DOMException('Aborted', 'AbortError'));
return;
}
const resolver: Resolver = { resolve, reject };
const list = pending.get(path);
if (list) {
list.push({ resolve, reject });
list.push(resolver);
} else {
pending.set(path, [{ resolve, reject }]);
pending.set(path, [resolver]);
}
if (signal) {
const onAbort = () => {
// Remove this resolver from the pending batch. If no other card is
// still waiting on the same path, drop the path entirely so the next
// batch doesn't include it.
const arr = pending.get(path);
if (arr) {
const idx = arr.indexOf(resolver);
if (idx >= 0) arr.splice(idx, 1);
if (arr.length === 0) pending.delete(path);
}
reject(new DOMException('Aborted', 'AbortError'));
};
signal.addEventListener('abort', onAbort, { once: true });
}
scheduleFlush();
});
@ -91,22 +111,24 @@ export default function useCaptionBatch(imgPath: string | null, refreshKey: numb
}
let cancelled = false;
const controller = new AbortController();
lastPathRef.current = imgPath;
setIsLoaded(false);
requestCaption(imgPath)
requestCaption(imgPath, controller.signal)
.then(value => {
if (cancelled || lastPathRef.current !== imgPath) return;
setCaption(value);
setIsLoaded(true);
})
.catch(err => {
if (cancelled) return;
if (err?.name === 'AbortError' || cancelled) return;
console.error('Error fetching caption:', err);
setIsLoaded(true);
});
return () => {
cancelled = true;
controller.abort();
};
}, [imgPath, refreshKey]);