From 7f9a142dfd2ce4a5b8f0600b382ccb1cce2416d7 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Wed, 12 Aug 2026 20:33:42 -0600 Subject: [PATCH] Generate thumbnails for dataset items so things load faster and the ui is more stable. Especially good for videos. --- .../captioner/BaseCaptioner.py | 3 +- toolkit/data_loader.py | 6 +- ui/cron/fileServer.ts | 150 +++++++++++++++++- ui/src/components/DatasetImageCard.tsx | 45 ++++-- 4 files changed, 185 insertions(+), 19 deletions(-) diff --git a/extensions_built_in/captioner/BaseCaptioner.py b/extensions_built_in/captioner/BaseCaptioner.py index 08724419..edfebb45 100644 --- a/extensions_built_in/captioner/BaseCaptioner.py +++ b/extensions_built_in/captioner/BaseCaptioner.py @@ -163,7 +163,8 @@ class BaseCaptioner(BaseExtensionProcess): def find_files(self): # recursivly find all the files in the path_to_caption with the specified extensions and save the paths to self.file_paths for root, dirs, files in os.walk(self.caption_config.path_to_caption): - dirs[:] = [d for d in dirs if d != "_controls"] + # skip _controls and hidden dirs (.thumbs, .tmp) + dirs[:] = [d for d in dirs if d != "_controls" and not d.startswith(".")] for file in files: if any( file.lower().endswith(f".{ext}") and not file.startswith(".") diff --git a/toolkit/data_loader.py b/toolkit/data_loader.py index 11dc879b..3e9773af 100644 --- a/toolkit/data_loader.py +++ b/toolkit/data_loader.py @@ -434,7 +434,11 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin # look for videos and images. Video models can train on both; # images are bucketed separately as single-frame items extensions = video_extensions + image_extensions - file_list = [os.path.join(root, file) for root, _, files in os.walk(self.dataset_path) for file in files if file.lower().endswith(tuple(extensions)) and not file.startswith('.')] + # prune hidden dirs (.thumbs, .tmp) so their contents never train + file_list = [] + for root, dirs, files in os.walk(self.dataset_path): + dirs[:] = [d for d in dirs if not d.startswith('.')] + file_list.extend(os.path.join(root, file) for file in files if file.lower().endswith(tuple(extensions)) and not file.startswith('.')) else: # assume json with open(self.dataset_path, 'r') as f: diff --git a/ui/cron/fileServer.ts b/ui/cron/fileServer.ts index 0f3276a2..128df58a 100644 --- a/ui/cron/fileServer.ts +++ b/ui/cron/fileServer.ts @@ -31,7 +31,7 @@ import os from 'os'; import path from 'path'; import { pipeline } from 'stream'; import prisma from './prisma'; -import { defaultDatasetsFolder, defaultTrainFolder, defaultDataRoot } from './paths'; +import { defaultDatasetsFolder, defaultTrainFolder, defaultDataRoot, TOOLKIT_ROOT } from './paths'; const isDev = process.argv.includes('dev'); @@ -83,6 +83,139 @@ async function getRoots(forceFresh = false): Promise { return roots; } +// --------------------------------------------------------------------------- +// Thumbnail generation for ?thumb=1 requests whose thumb doesn't exist yet. +// Output matches the Python generator (SampleConfig._generate_thumbnail in +// toolkit/config_modules.py): 300x300 center-cropped q90 jpg written +// atomically into the sibling .thumbs folder as ..jpg. +// --------------------------------------------------------------------------- +const THUMB_SIZE = 300; +const IMAGE_THUMB_EXTS = new Set(['.jpg', '.jpeg', '.png', '.webp', '.gif', '.bmp']); +const VIDEO_THUMB_EXTS = new Set(['.mp4', '.avi', '.mov', '.mkv', '.wmv', '.m4v', '.flv']); + +// sharp is not a direct dependency; Next.js vendors it (for next/image), so +// resolve it out of next's node_modules tree. If that ever fails, image +// thumbs just fall back to serving the original file. +const sharp: any = (() => { + try { + return require(require.resolve('sharp', { paths: [path.dirname(require.resolve('next/package.json'))] })); + } catch { + return null; + } +})(); + +// The manager provisions a portable FFmpeg at /.ffmpeg (see +// manager/ffmpeg.py) — prefer it over whatever is on PATH. Its Linux build is +// a shared one, so spawning it directly (i.e. not via `manager launch`, which +// sets this up itself) needs .ffmpeg/lib on LD_LIBRARY_PATH. +const localFfmpegExe = path.join(TOOLKIT_ROOT, '.ffmpeg', 'bin', process.platform === 'win32' ? 'ffmpeg.exe' : 'ffmpeg'); +const ffmpegExe = fs.existsSync(localFfmpegExe) ? localFfmpegExe : 'ffmpeg'; +const ffmpegEnv: NodeJS.ProcessEnv = (() => { + const libDir = path.join(TOOLKIT_ROOT, '.ffmpeg', 'lib'); + if (ffmpegExe === 'ffmpeg' || process.platform !== 'linux' || !fs.existsSync(libDir)) return process.env; + const prior = process.env.LD_LIBRARY_PATH; + return { ...process.env, LD_LIBRARY_PATH: prior ? `${libDir}${path.delimiter}${prior}` : libDir }; +})(); + +let ffmpegMissing = false; +// A generation attempt that failed is memoized (keyed on source mtime, so a +// re-written file retries) — the gallery re-requests thumbs constantly and +// must not re-run ffmpeg/sharp against a broken file on every poll. +const failedThumbs = new Set(); +const inFlightThumbs = new Map>(); + +// A gallery burst can request hundreds of missing thumbs at once; cap how +// many decodes/ffmpeg spawns run concurrently per worker. +const MAX_THUMB_GEN = 4; +let activeThumbGen = 0; +const thumbGenQueue: (() => void)[] = []; +async function withThumbGenSlot(fn: () => Promise): Promise { + if (activeThumbGen >= MAX_THUMB_GEN) { + await new Promise(r => thumbGenQueue.push(r)); + } + activeThumbGen++; + try { + return await fn(); + } finally { + activeThumbGen--; + thumbGenQueue.shift()?.(); + } +} + +async function generateThumb(sourcePath: string, thumbPath: string): Promise { + const ext = path.extname(sourcePath).toLowerCase(); + const isImage = IMAGE_THUMB_EXTS.has(ext); + const isVideo = VIDEO_THUMB_EXTS.has(ext); + if ((isImage && !sharp) || (isVideo && ffmpegMissing) || (!isImage && !isVideo)) return false; + await fs.promises.mkdir(path.dirname(thumbPath), { recursive: true }); + // Write to a per-process tmp name, then atomically rename into place (same + // as the Python generator) so a concurrent request never reads a partial + // thumb. The .jpg suffix is required for ffmpeg's output format detection. + const tmpPath = `${thumbPath}.${process.pid}.tmp.jpg`; + try { + if (isImage) { + // sharp opens animated formats on the first frame by default + await sharp(sourcePath) + .resize(THUMB_SIZE, THUMB_SIZE, { fit: 'cover' }) + .jpeg({ quality: 90 }) + .toFile(tmpPath); + } else { + await new Promise((resolve, reject) => { + const child = spawn( + ffmpegExe, + [ + '-y', + '-loglevel', 'error', + '-i', sourcePath, + '-frames:v', '1', + '-vf', `crop='min(iw,ih)':'min(iw,ih)',scale=${THUMB_SIZE}:${THUMB_SIZE}`, + '-q:v', '2', + tmpPath, + ], + { stdio: ['ignore', 'ignore', 'pipe'], env: ffmpegEnv }, + ); + let stderr = ''; + child.stderr!.on('data', chunk => (stderr += chunk.toString())); + const timer = setTimeout(() => child.kill('SIGKILL'), 30_000); + child.on('error', (err: NodeJS.ErrnoException) => { + clearTimeout(timer); + if (err.code === 'ENOENT') ffmpegMissing = true; + reject(err); + }); + child.on('exit', code => { + clearTimeout(timer); + code === 0 ? resolve() : reject(new Error(`ffmpeg exited with ${code}: ${stderr.trim()}`)); + }); + }); + } + await fs.promises.rename(tmpPath, thumbPath); + return true; + } catch (err) { + await fs.promises.unlink(tmpPath).catch(() => { }); + throw err; + } +} + +function ensureThumb(sourcePath: string, thumbPath: string, sourceMtimeMs: number): Promise { + const failKey = `${thumbPath}:${sourceMtimeMs}`; + if (failedThumbs.has(failKey)) return Promise.resolve(false); + let pending = inFlightThumbs.get(thumbPath); + if (!pending) { + pending = withThumbGenSlot(() => generateThumb(sourcePath, thumbPath)) + .catch(err => { + console.warn(`Failed to generate thumbnail for ${sourcePath}: ${err?.message || err}`); + return false; + }) + .then(ok => { + if (!ok) failedThumbs.add(failKey); + return ok; + }) + .finally(() => inFlightThumbs.delete(thumbPath)); + inFlightThumbs.set(thumbPath, pending); + } + return pending; +} + const contentTypeMap: { [key: string]: string } = { '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', @@ -143,12 +276,19 @@ async function serveFile(req: http.IncomingMessage, res: http.ServerResponse, pr return; } - // ?thumb=1 serves the pre-generated 300x300 jpg from the sibling .thumbs - // folder (..jpg) when it exists; otherwise falls through to - // the full file exactly as before. Mirrors the Next.js /api/img route. + // ?thumb=1 serves the 300x300 jpg from the sibling .thumbs folder + // (..jpg), generating and saving it on the fly when missing. + // Falls through to the full file only if generation isn't possible + // (unsupported format, no ffmpeg, corrupt file). 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); + let thumbStat = await fs.promises.stat(thumbPath).catch(() => null); + if (!(thumbStat && thumbStat.isFile())) { + const srcStat = await fs.promises.stat(resolvedFilePath).catch(() => null); + if (srcStat && srcStat.isFile() && (await ensureThumb(resolvedFilePath, thumbPath, srcStat.mtimeMs))) { + thumbStat = await fs.promises.stat(thumbPath).catch(() => null); + } + } if (thumbStat && thumbStat.isFile()) { resolvedFilePath = thumbPath; } diff --git a/ui/src/components/DatasetImageCard.tsx b/ui/src/components/DatasetImageCard.tsx index 5d280d64..2f8f483f 100644 --- a/ui/src/components/DatasetImageCard.tsx +++ b/ui/src/components/DatasetImageCard.tsx @@ -1,5 +1,5 @@ import React, { useEffect, useState, ReactNode, KeyboardEvent, useRef } from 'react'; -import { FaTrashAlt } from 'react-icons/fa'; +import { FaTrashAlt, FaPlay } from 'react-icons/fa'; import { openConfirm } from './ConfirmModal'; import classNames from 'classnames'; import { apiClient } from '@/utils/api'; @@ -38,12 +38,12 @@ const DatasetImageCard: React.FC = ({ const [showAudioPlayer, setShowAudioPlayer] = useState(true); const [pollTick, setPollTick] = useState(0); const [blobUrl, setBlobUrl] = useState(null); + const [streamVideo, setStreamVideo] = useState(false); const [isVisible, setIsVisible] = useState(false); const cardRef = useRef(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. @@ -72,9 +72,12 @@ const DatasetImageCard: React.FC = ({ // Drive image loads through fetch + AbortController so scrolling past actually // cancels in-flight requests. Debounced 80ms so fast scroll-throughs never - // start a request. + // start a request. Both images and videos pull the 300x300 thumb (the server + // generates it on a miss); ?thumb=1 falls through to the real file only when + // a thumb can't be made, so a video/* response means "no thumb available" — + // abort before downloading the body and stream a