Generate thumbnails for dataset items so things load faster and the ui is more stable. Especially good for videos.

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(".")

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:

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<Roots> {
   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 <name>.<ext>.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 <repo>/.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<string>();
+const inFlightThumbs = new Map<string, Promise<boolean>>();
+
+// 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<T>(fn: () => Promise<T>): Promise<T> {
+  if (activeThumbGen >= MAX_THUMB_GEN) {
+    await new Promise<void>(r => thumbGenQueue.push(r));
+  }
+  activeThumbGen++;
+  try {
+    return await fn();
+  } finally {
+    activeThumbGen--;
+    thumbGenQueue.shift()?.();
+  }
+}
+
+async function generateThumb(sourcePath: string, thumbPath: string): Promise<boolean> {
+  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<void>((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<boolean> {
+  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
+    // ?thumb=1 serves the 300x300 jpg from the sibling .thumbs folder
-    // folder (<name>.<ext>.jpg) when it exists; otherwise falls through to
+    // (<name>.<ext>.jpg), generating and saving it on the fly when missing.
-    // the full file exactly as before. Mirrors the Next.js /api/img route.
+    // 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;
       }

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<DatasetImageCardProps> = ({
   const [showAudioPlayer, setShowAudioPlayer] = useState(true);
   const [pollTick, setPollTick] = useState(0);
   const [blobUrl, setBlobUrl] = useState<string | null>(null);
+  const [streamVideo, setStreamVideo] = useState(false);
   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.
@@ -72,9 +72,12 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
 
   // 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 <video> tag instead.
   useEffect(() => {
-    if (!isItImage) return;
+    if (isItAudio) return;
     if (!isVisible) return;
 
     const controller = new AbortController();
@@ -82,13 +85,21 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
     let objectUrl: string | null = null;
 
     const timer = window.setTimeout(() => {
-      fetch(`/api/img/${encodeURIComponent(imageUrl)}`, { signal: controller.signal })
+      fetch(`/api/img/${encodeURIComponent(imageUrl)}?thumb=1`, { signal: controller.signal })
         .then(r => {
           if (!r.ok) throw new Error(`HTTP ${r.status}`);
+          if ((r.headers.get('content-type') || '').startsWith('video/')) {
+            controller.abort();
+            if (!cancelled) {
+              setStreamVideo(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);
@@ -104,9 +115,10 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
       controller.abort();
       if (objectUrl) URL.revokeObjectURL(objectUrl);
       setBlobUrl(null);
+      setStreamVideo(false);
       setLoaded(false);
     };
-  }, [imageUrl, isItImage, isVisible]);
+  }, [imageUrl, isItAudio, isVisible]);
 
   const combinedRefreshKey = captionRefreshKey + pollTick;
   const { caption: fetchedCaption, isLoaded: isCaptionLoaded } = useCaptionBatch(
@@ -195,17 +207,21 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
       <div className="relative w-full" style={{ paddingBottom: '100%' }}>
         <div
           className={classNames('absolute inset-0 rounded-t-lg shadow-md bg-gray-900', {
-            'animate-pulse': isItImage && !loaded,
+            'animate-pulse': !isItAudio && !loaded,
           })}
         >
-          {isItAVideo && (
+          {streamVideo && (
             <video
               src={`/api/img/${encodeURIComponent(imageUrl)}`}
-              className={`w-full h-full object-contain`}
+              className={classNames('w-full h-full object-contain', {
+                'cursor-zoom-in': !!onImageClick,
+              })}
+              onClick={onImageClick}
               autoPlay={false}
+              preload="metadata"
+              playsInline
               loop
               muted
-              controls
             />
           )}
           {isItAudio && !showAudioPlayer && (
@@ -226,7 +242,7 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
           {isItAudio && showAudioPlayer && (
             <AudioPlayer src={`/api/img/${encodeURIComponent(imageUrl)}`} title={imageUrl.replace(/^.*[\\/]/, '')} />
           )}
-          {isItImage && blobUrl && (
+          {!isItAudio && blobUrl && (
             <img
               src={blobUrl}
               alt={alt}
@@ -236,6 +252,11 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
               })}
             />
           )}
+          {isItAVideo && loaded && (
+            <div className="absolute bottom-2 left-2 bg-gray-900/70 rounded-full p-2 pointer-events-none">
+              <FaPlay className="w-3 h-3 text-white" />
+            </div>
+          )}
           {children && <div className="absolute inset-0 flex items-center justify-center">{children}</div>}
           <div className="absolute top-1 right-1 flex space-x-2 z-10">
             <button