Drastically improved the performance of the dataset viewer on large datasets by switching to virtualization. Fixed issue with images loading when they have two periods in a row ..

2026-05-24 15:13:55 -06:00 · 2026-05-24 15:13:55 -06:00 · 307ff11bc5
parent c6a7e81a70
commit 307ff11bc5
12 changed files with 360 additions and 219 deletions
--- a/ui/package-lock.json
+++ b/ui/package-lock.json
@ -24,6 +24,7 @@
        "react-global-hooks": "^1.3.5",
        "react-icons": "^5.5.0",
        "react-select": "^5.10.1",
+        "react-virtuoso": "^4.18.7",
        "react-zoom-pan-pinch": "^4.0.3",
        "sqlite3": "^5.1.7",
        "systeminformation": "^5.27.11",
@ -5167,6 +5168,16 @@
        "react-dom": ">=16.6.0"
      }
    },
+    "node_modules/react-virtuoso": {
+      "version": "4.18.7",
+      "resolved": "https://registry.npmjs.org/react-virtuoso/-/react-virtuoso-4.18.7.tgz",
+      "integrity": "sha512-xNF5zDGEEIMB7cKwcen/pLig0YDf6OnfFrVgKFa7sHPf9fRem0CaLshyObbBcP88jzn0enavL39EgplgdyT21g==",
+      "license": "MIT",
+      "peerDependencies": {
+        "react": ">=16 || >=17 || >= 18 || >= 19",
+        "react-dom": ">=16 || >=17 || >= 18 || >=19"
+      }
+    },
    "node_modules/react-zoom-pan-pinch": {
      "version": "4.0.3",
      "resolved": "https://registry.npmjs.org/react-zoom-pan-pinch/-/react-zoom-pan-pinch-4.0.3.tgz",
--- a/ui/package.json
+++ b/ui/package.json
@ -28,6 +28,7 @@
    "react-global-hooks": "^1.3.5",
    "react-icons": "^5.5.0",
    "react-select": "^5.10.1",
+    "react-virtuoso": "^4.18.7",
    "react-zoom-pan-pinch": "^4.0.3",
    "sqlite3": "^5.1.7",
    "systeminformation": "^5.27.11",
--- a/ui/src/app/api/audio/art/[...audioPath]/route.ts
+++ b/ui/src/app/api/audio/art/[...audioPath]/route.ts
@ -94,9 +94,9 @@ function extractArtFromTag(buf: Buffer): ArtResult {
        verMajor === 4
          ? synchsafeToInt(tagData[offset + 4], tagData[offset + 5], tagData[offset + 6], tagData[offset + 7])
          : (tagData[offset + 4] << 24) |
-            (tagData[offset + 5] << 16) |
-            (tagData[offset + 6] << 8) |
-            tagData[offset + 7];
+          (tagData[offset + 5] << 16) |
+          (tagData[offset + 6] << 8) |
+          tagData[offset + 7];
      const flag2 = tagData[offset + 9];
      offset += 10;
      if (!id.trim() || size <= 0 || offset + size > tagData.length) break;
@ -135,19 +135,22 @@ export async function GET(request: NextRequest, { params }: { params: { audioPat
    const trainingRoot = await getTrainingFolder();
    const dataRoot = await getDataRoot();
    const allowedDirs = [datasetRoot, trainingRoot, dataRoot];
-    const isAllowed = allowedDirs.some(d => filepath.startsWith(d)) && !filepath.includes('..');
+    // Resolve so `..` segments collapse, then verify still under an allowed root.
+    // Substring `.includes('..')` false-positives on filenames containing `..` as text.
+    const resolved = path.resolve(filepath);
+    const isAllowed = allowedDirs.some(d => resolved === d || resolved.startsWith(d + path.sep));
    if (!isAllowed) {
      return new NextResponse('Access denied', { status: 403 });
    }

-    const stat = await fs.promises.stat(filepath).catch(() => null);
+    const stat = await fs.promises.stat(resolved).catch(() => null);
    if (!stat || !stat.isFile()) {
      return new NextResponse('File not found', { status: 404 });
    }

    // Read only the ID3 tag (first min(tagSize, 4MB) bytes).
    // First read 10 bytes to get tag size, then read the full tag.
-    const fd = await fs.promises.open(filepath, 'r');
+    const fd = await fs.promises.open(resolved, 'r');
    try {
      const headerBuf = Buffer.alloc(10);
      await fd.read(headerBuf, 0, 10, 0);
@ -167,7 +170,7 @@ export async function GET(request: NextRequest, { params }: { params: { audioPat
        return new NextResponse('No album art found', { status: 404 });
      }

-      return new NextResponse(art.data, {
+      return new NextResponse(art.data as any, {
        headers: {
          'Content-Type': art.mime,
          'Content-Length': String(art.data.length),
--- a/ui/src/app/api/caption/get/route.ts
+++ b/ui/src/app/api/caption/get/route.ts
@ -4,6 +4,11 @@ import fs from 'fs';
 import path from 'path';
 import { getDatasetsRoot } from '@/server/settings';

+function isUnderRoot(filepath: string, root: string): boolean {
+  const resolved = path.resolve(filepath);
+  return resolved === root || resolved.startsWith(root + path.sep);
+}
+
 export async function POST(request: NextRequest) {
  let body;
  try {
@ -30,8 +35,10 @@ export async function POST(request: NextRequest) {
    // Get allowed directories
    const allowedDir = await getDatasetsRoot();

-    // Security check: Ensure path is in allowed directory
-    const isAllowed = filepath.startsWith(allowedDir) && !filepath.includes('..');
+    // Security check: resolve so `..` segments collapse, then verify it's still
+    // under the allowed root. Substring `.includes('..')` would false-positive
+    // on filenames that contain `..` as text (e.g. an ellipsis in a filename).
+    const isAllowed = isUnderRoot(filepath, allowedDir);

    if (!isAllowed) {
      console.warn(`Access denied: ${filepath} not in ${allowedDir}`);
--- a/ui/src/app/api/caption/getBatch/route.ts
+++ b/ui/src/app/api/caption/getBatch/route.ts
@ -0,0 +1,45 @@
+/* eslint-disable */
+import { NextRequest, NextResponse } from 'next/server';
+import fs from 'fs';
+import path from 'path';
+import { getDatasetsRoot } from '@/server/settings';
+
+function isUnderRoot(filepath: string, root: string): boolean {
+  const resolved = path.resolve(filepath);
+  return resolved === root || resolved.startsWith(root + path.sep);
+}
+
+export async function POST(request: NextRequest) {
+  let body;
+  try {
+    body = await request.json();
+  } catch {
+    return new NextResponse(null, { status: 499 });
+  }
+
+  if (request.signal.aborted) {
+    return new NextResponse(null, { status: 499 });
+  }
+
+  const { imgPaths } = body as { imgPaths?: string[] };
+  if (!Array.isArray(imgPaths)) {
+    return NextResponse.json({ error: 'imgPaths must be an array' }, { status: 400 });
+  }
+
+  const allowedDir = await getDatasetsRoot();
+  const captions: Record<string, string> = {};
+
+  for (const imgPath of imgPaths) {
+    if (typeof imgPath !== 'string') continue;
+    if (!isUnderRoot(imgPath, allowedDir)) continue;
+
+    const captionPath = imgPath.replace(/\.[^/.]+$/, '') + '.txt';
+    try {
+      captions[imgPath] = fs.existsSync(captionPath) ? fs.readFileSync(captionPath, 'utf-8') : '';
+    } catch {
+      captions[imgPath] = '';
+    }
+  }
+
+  return NextResponse.json({ captions });
+}
--- a/ui/src/app/api/datasets/listImages/route.ts
+++ b/ui/src/app/api/datasets/listImages/route.ts
@ -18,6 +18,9 @@ export async function POST(request: Request) {
    // Find all images recursively
    const imageFiles = findImagesRecursively(datasetFolder);

+    // Sort server-side so the client doesn't have to sort large lists
+    imageFiles.sort((a, b) => a.localeCompare(b));
+
    // Format response
    const result = imageFiles.map(imgPath => ({
      img_path: imgPath,
@ -39,19 +42,20 @@ function findImagesRecursively(dir: string): string[] {
  const imageExtensions = ['.png', '.jpg', '.jpeg', '.webp', '.mp4', '.avi', '.mov', '.mkv', '.wmv', '.m4v', '.flv', '.mp3', '.wav', '.flac', '.ogg'];
  let results: string[] = [];

-  const items = fs.readdirSync(dir);
+  // withFileTypes avoids a separate statSync per entry — a big win on large datasets
+  const entries = fs.readdirSync(dir, { withFileTypes: true });

-  for (const item of items) {
-    const itemPath = path.join(dir, item);
-    const stat = fs.statSync(itemPath);
+  for (const entry of entries) {
+    const name = entry.name;
+    if (name.startsWith('.')) continue;
+    const itemPath = path.join(dir, name);

-    if (stat.isDirectory() && item !== '_controls' && !item.startsWith('.')) {
-      // If it's a directory, recursively search it
+    if (entry.isDirectory()) {
+      if (name === '_controls') continue;
      results = results.concat(findImagesRecursively(itemPath));
-    } else {
-      // If it's a file, check if it's an image
-      const ext = path.extname(itemPath).toLowerCase();
-      if (imageExtensions.includes(ext) && !item.startsWith('.')) {
+    } else if (entry.isFile()) {
+      const ext = path.extname(name).toLowerCase();
+      if (imageExtensions.includes(ext)) {
        results.push(itemPath);
      }
    }
--- a/ui/src/app/api/files/[...filePath]/route.ts
+++ b/ui/src/app/api/files/[...filePath]/route.ts
@ -15,32 +15,36 @@ export async function GET(request: NextRequest, { params }: { params: { filePath
    const trainingRoot = await getTrainingFolder();
    const allowedDirs = [datasetRoot, trainingRoot];

-    // Security check: Ensure path is in allowed directory
-    const isAllowed =
-      allowedDirs.some(allowedDir => decodedFilePath.startsWith(allowedDir)) && !decodedFilePath.includes('..');
+    // Security check: resolve so `..` segments collapse, then verify still under
+    // an allowed root. Substring `.includes('..')` false-positives on filenames
+    // containing `..` as text (e.g. an ellipsis in a filename).
+    const resolvedFilePath = path.resolve(decodedFilePath);
+    const isAllowed = allowedDirs.some(
+      allowedDir => resolvedFilePath === allowedDir || resolvedFilePath.startsWith(allowedDir + path.sep),
+    );

    if (!isAllowed) {
-      console.warn(`Access denied: ${decodedFilePath} not in ${allowedDirs.join(', ')}`);
+      console.warn(`Access denied: ${resolvedFilePath} not in ${allowedDirs.join(', ')}`);
      return new NextResponse('Access denied', { status: 403 });
    }

    // Check if file exists
-    if (!fs.existsSync(decodedFilePath)) {
-      console.warn(`File not found: ${decodedFilePath}`);
+    if (!fs.existsSync(resolvedFilePath)) {
+      console.warn(`File not found: ${resolvedFilePath}`);
      return new NextResponse('File not found', { status: 404 });
    }

    // Get file info
-    const stat = fs.statSync(decodedFilePath);
+    const stat = fs.statSync(resolvedFilePath);
    if (!stat.isFile()) {
      return new NextResponse('Not a file', { status: 400 });
    }

    // Get filename for Content-Disposition
-    const filename = path.basename(decodedFilePath);
+    const filename = path.basename(resolvedFilePath);

    // Determine content type
-    const ext = path.extname(decodedFilePath).toLowerCase();
+    const ext = path.extname(resolvedFilePath).toLowerCase();
    const contentTypeMap: { [key: string]: string } = {
      '.jpg': 'image/jpeg',
      '.jpeg': 'image/jpeg',
@ -87,7 +91,7 @@ export async function GET(request: NextRequest, { params }: { params: { filePath
      const end = parts[1] ? parseInt(parts[1], 10) : Math.min(start + 10 * 1024 * 1024, stat.size - 1); // 10MB chunks
      const chunkSize = end - start + 1;

-      const fileStream = fs.createReadStream(decodedFilePath, {
+      const fileStream = fs.createReadStream(resolvedFilePath, {
        start,
        end,
        highWaterMark: 64 * 1024, // 64KB buffer
@ -103,7 +107,7 @@ export async function GET(request: NextRequest, { params }: { params: { filePath
      });
    } else {
      // For full file download, read directly without streaming wrapper
-      const fileStream = fs.createReadStream(decodedFilePath, {
+      const fileStream = fs.createReadStream(resolvedFilePath, {
        highWaterMark: 64 * 1024, // 64KB buffer
      });

--- a/ui/src/app/api/img/[...imagePath]/route.ts
+++ b/ui/src/app/api/img/[...imagePath]/route.ts
@ -41,21 +41,26 @@ export async function GET(request: NextRequest, { params }: { params: { imagePat

    const allowedDirs = [datasetRoot, trainingRoot, dataRoot];

-    // Security check: Ensure path is in allowed directory
-    const isAllowed = allowedDirs.some(allowedDir => filepath.startsWith(allowedDir)) && !filepath.includes('..');
+    // 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);
+    const isAllowed = allowedDirs.some(
+      allowedDir => resolved === allowedDir || resolved.startsWith(allowedDir + path.sep),
+    );

    if (!isAllowed) {
-      console.warn(`Access denied: ${filepath} not in ${allowedDirs.join(', ')}`);
+      console.warn(`Access denied: ${resolved} not in ${allowedDirs.join(', ')}`);
      return new NextResponse('Access denied', { status: 403 });
    }

    // Stat file (async)
-    const stat = await fs.promises.stat(filepath).catch(() => null);
+    const stat = await fs.promises.stat(resolved).catch(() => null);
    if (!stat || !stat.isFile()) {
      return new NextResponse('File not found', { status: 404 });
    }

-    const ext = path.extname(filepath).toLowerCase();
+    const ext = path.extname(resolved).toLowerCase();
    const contentType = contentTypeMap[ext] || 'application/octet-stream';

    // Support range requests for video/audio seeking
@ -66,7 +71,7 @@ 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(filepath, { start, end });
+      const stream = fs.createReadStream(resolved, { start, end });
      const readable = new ReadableStream({
        start(controller) {
          stream.on('data', chunk => controller.enqueue(chunk));
@ -91,7 +96,7 @@ export async function GET(request: NextRequest, { params }: { params: { imagePat
    }

    // Stream the file instead of buffering it entirely
-    const stream = fs.createReadStream(filepath);
+    const stream = fs.createReadStream(resolved);
    const readable = new ReadableStream({
      start(controller) {
        stream.on('data', chunk => controller.enqueue(chunk));
--- a/ui/src/app/datasets/[datasetName]/page.tsx
+++ b/ui/src/app/datasets/[datasetName]/page.tsx
@ -1,8 +1,9 @@
 'use client';

-import { useEffect, useState, use, useMemo } from 'react';
+import { useEffect, useState, use, useMemo, useCallback } from 'react';
 import { LuImageOff, LuLoader, LuBan } from 'react-icons/lu';
 import { FaChevronLeft } from 'react-icons/fa';
+import { VirtuosoGrid } from 'react-virtuoso';
 import DatasetImageCard from '@/components/DatasetImageCard';
 import DatasetImageViewer from '@/components/DatasetImageViewer';
 import { Button } from '@headlessui/react';
@ -22,17 +23,16 @@ export default function DatasetPage({ params }: { params: { datasetName: string
  const { settings, isSettingsLoaded } = useSettings();
  const [selectedImgPath, setSelectedImgPath] = useState<string | null>(null);
  const [captionRefreshKeys, setCaptionRefreshKeys] = useState<Record<string, number>>({});
+  const [scrollParent, setScrollParent] = useState<HTMLDivElement | null>(null);
+  const scrollParentCallback = useCallback((el: HTMLDivElement | null) => setScrollParent(el), []);

  const refreshImageList = (dbName: string) => {
    setStatus('loading');
-    console.log('Fetching images for dataset:', dbName);
    apiClient
      .post('/api/datasets/listImages', { datasetName: dbName })
      .then((res: any) => {
        const data = res.data;
-        console.log('Images:', data.images);
-        // sort
-        data.images.sort((a: { img_path: string }, b: { img_path: string }) => a.img_path.localeCompare(b.img_path));
+        // Server already sorts; avoid the client-side sort that's expensive on large lists.
        setImgList(data.images);
        setStatus('success');
      })
@ -43,6 +43,8 @@ export default function DatasetPage({ params }: { params: { datasetName: string
  };
  useOpenImagesModalOnDrag(datasetName, () => refreshImageList(datasetName));

+  const imgPaths = useMemo(() => imgList.map(img => img.img_path), [imgList]);
+
  useEffect(() => {
    if (datasetName) {
      refreshImageList(datasetName);
@ -129,28 +131,36 @@ export default function DatasetPage({ params }: { params: { datasetName: string
          </Button>
        </div>
      </TopBar>
-      <MainContent>
+      <MainContent ref={scrollParentCallback}>
        {PageInfoContent}
-        {status === 'success' && imgList.length > 0 && (
-          <div className="grid grid-cols-1 sm:grid-cols-2 md:grid-cols-3 lg:grid-cols-4 gap-4">
-            {imgList.map(img => (
-              <DatasetImageCard
-                key={img.img_path}
-                alt="image"
-                isAutoCaptioning={isAutoCaptioning}
-                imageUrl={img.img_path}
-                onDelete={() => refreshImageList(datasetName)}
-                onImageClick={() => setSelectedImgPath(img.img_path)}
-                captionRefreshKey={captionRefreshKeys[img.img_path] || 0}
-              />
-            ))}
-          </div>
+        {status === 'success' && imgList.length > 0 && scrollParent && (
+          <VirtuosoGrid
+            totalCount={imgList.length}
+            customScrollParent={scrollParent}
+            overscan={400}
+            listClassName="grid grid-cols-1 sm:grid-cols-2 md:grid-cols-3 lg:grid-cols-4 gap-4"
+            itemContent={index => {
+              const img = imgList[index];
+              if (!img) return null;
+              return (
+                <DatasetImageCard
+                  alt="image"
+                  isAutoCaptioning={isAutoCaptioning}
+                  imageUrl={img.img_path}
+                  onDelete={() => refreshImageList(datasetName)}
+                  onImageClick={() => setSelectedImgPath(img.img_path)}
+                  captionRefreshKey={captionRefreshKeys[img.img_path] || 0}
+                />
+              );
+            }}
+            computeItemKey={index => imgList[index]?.img_path ?? index}
+          />
        )}
      </MainContent>
      <AddImagesModal />
      <DatasetImageViewer
        imgPath={selectedImgPath}
-        imageList={imgList.map(img => img.img_path)}
+        imageList={imgPaths}
        onChange={setSelectedImgPath}
        refreshImages={() => refreshImageList(datasetName)}
        onCaptionSaved={path => setCaptionRefreshKeys(prev => ({ ...prev, [path]: (prev[path] || 0) + 1 }))}
--- a/ui/src/components/DatasetImageCard.tsx
+++ b/ui/src/components/DatasetImageCard.tsx
@ -1,10 +1,11 @@
-import React, { useRef, useEffect, useState, ReactNode, KeyboardEvent } from 'react';
-import { FaTrashAlt, FaEye, FaEyeSlash } from 'react-icons/fa';
+import React, { useEffect, useState, ReactNode, KeyboardEvent, useRef } from 'react';
+import { FaTrashAlt } from 'react-icons/fa';
 import { openConfirm } from './ConfirmModal';
 import classNames from 'classnames';
 import { apiClient } from '@/utils/api';
 import AudioPlayer from './AudioPlayer';
 import { isVideo, isAudio } from '@/utils/basic';
+import useCaptionBatch, { setCachedCaption } from '@/hooks/useCaptionBatch';

 interface DatasetImageCardProps {
  imageUrl: string;
@ -27,131 +28,85 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
  onImageClick,
  captionRefreshKey = 0,
 }) => {
-  const cardRef = useRef<HTMLDivElement>(null);
-  const [isVisible, setIsVisible] = useState<boolean>(false);
-  const [inViewport, setInViewport] = useState<boolean>(false);
  const [loaded, setLoaded] = useState<boolean>(false);
-  const [isCaptionLoaded, setIsCaptionLoaded] = useState<boolean>(false);
+  const [showAudioPlayer, setShowAudioPlayer] = useState(true);
+  const [pollTick, setPollTick] = useState(0);
+
+  const combinedRefreshKey = captionRefreshKey + pollTick;
+  const { caption: fetchedCaption, isLoaded: isCaptionLoaded } = useCaptionBatch(imageUrl, combinedRefreshKey);
+
  const [caption, setCaption] = useState<string>('');
  const [savedCaption, setSavedCaption] = useState<string>('');
-  const abortControllerRef = useRef<AbortController | null>(null);
+  const dirtyRef = useRef<boolean>(false);

-  const fetchCaption = async () => {
-    if (isCaptionLoaded) return;
-    abortControllerRef.current?.abort();
-    const controller = new AbortController();
-    abortControllerRef.current = controller;
-    apiClient
-      .post(`/api/caption/get`, { imgPath: imageUrl }, { signal: controller.signal })
-      .then(res => res.data)
-      .then(data => {
-        console.log('Caption fetched:', data);
-        if (data) {
-          data = `${data}`;
-        }
-        setCaption(data || '');
-        setSavedCaption((data || '').trim());
-        setIsCaptionLoaded(true);
-      })
-      .catch(error => {
-        if (controller.signal.aborted) return;
-        console.error('Error fetching caption:', error);
-      })
-      .finally(() => {
-        if (abortControllerRef.current === controller) {
-          abortControllerRef.current = null;
-        }
-      });
-  };
+  // Sync from the fetched caption, but don't clobber unsaved local edits.
+  useEffect(() => {
+    if (!isCaptionLoaded) return;
+    if (dirtyRef.current) return;
+    setCaption(fetchedCaption);
+    setSavedCaption(fetchedCaption.trim());
+  }, [fetchedCaption, isCaptionLoaded]);
+
+  // Poll while auto-captioning so backend-written captions show up.
+  useEffect(() => {
+    if (!isAutoCaptioning) return;
+    const interval = setInterval(() => setPollTick(t => t + 1), 5000);
+    return () => clearInterval(interval);
+  }, [isAutoCaptioning]);

  const saveCaption = () => {
    const trimmedCaption = caption.trim();
-    if (trimmedCaption === savedCaption) return;
+    if (trimmedCaption === savedCaption) {
+      dirtyRef.current = false;
+      return;
+    }
    apiClient
      .post('/api/img/caption', { imgPath: imageUrl, caption: trimmedCaption })
-      .then(res => res.data)
-      .then(data => {
-        console.log('Caption saved:', data);
+      .then(() => {
        setSavedCaption(trimmedCaption);
+        setCachedCaption(imageUrl, trimmedCaption);
+        dirtyRef.current = false;
      })
      .catch(error => {
        console.error('Error saving caption:', error);
      });
  };

-  // Only fetch caption when the component is both in viewport and visible
+  // Save any pending edit if the card unmounts (e.g. scrolled out of the virtualized window).
+  const latestRef = useRef({ caption, savedCaption, imageUrl });
  useEffect(() => {
-    if (inViewport && isVisible) {
-      fetchCaption();
-    }
-  }, [inViewport, isVisible, isCaptionLoaded]);
-
-  // Poll for caption updates every 5 seconds while auto-captioning
+    latestRef.current = { caption, savedCaption, imageUrl };
+  });
  useEffect(() => {
-    if (!isAutoCaptioning || !inViewport || !isVisible) return;
-    const interval = setInterval(() => {
-      // Reset so fetchCaption will re-fetch
-      setIsCaptionLoaded(false);
-    }, 5000);
-    return () => clearInterval(interval);
-  }, [isAutoCaptioning, inViewport, isVisible]);
-
-  // External trigger (e.g. caption edited in the full-screen viewer) — re-fetch
-  useEffect(() => {
-    if (captionRefreshKey === 0) return;
-    setIsCaptionLoaded(false);
-  }, [captionRefreshKey]);
-
-  useEffect(() => {
-    // Create intersection observer to check viewport visibility
-    const observer = new IntersectionObserver(
-      entries => {
-        if (entries[0].isIntersecting) {
-          setInViewport(true);
-          // Initialize isVisible to true when first coming into view
-          if (!isVisible) {
-            setIsVisible(true);
-          }
-        } else {
-          setInViewport(false);
-          // Cancel any in-flight caption fetch when scrolling away
-          abortControllerRef.current?.abort();
-        }
-      },
-      { threshold: 0.1 },
-    );
-
-    if (cardRef.current) {
-      observer.observe(cardRef.current);
-    }
-
    return () => {
-      observer.disconnect();
+      if (!dirtyRef.current) return;
+      const { caption: c, savedCaption: s, imageUrl: url } = latestRef.current;
+      const trimmed = c.trim();
+      if (trimmed === s) return;
+      apiClient
+        .post('/api/img/caption', { imgPath: url, caption: trimmed })
+        .then(() => setCachedCaption(url, trimmed))
+        .catch(err => console.error('Error saving caption on unmount:', err));
    };
  }, []);

-  const toggleVisibility = (): void => {
-    setIsVisible(prev => !prev);
-    if (!isVisible && !isCaptionLoaded) {
-      fetchCaption();
-    }
-  };
-
  const handleLoad = (): void => {
    setLoaded(true);
  };

  const handleKeyDown = (e: KeyboardEvent<HTMLTextAreaElement>): void => {
-    // If Enter is pressed without Shift, prevent default behavior and save
    if (e.key === 'Enter' && !e.shiftKey) {
      e.preventDefault();
      saveCaption();
    }
  };

-  const isCaptionCurrent = caption.trim() === savedCaption;
+  const handleCaptionChange = (value: string) => {
+    dirtyRef.current = value.trim() !== savedCaption;
+    setCaption(value);
+  };

-  const [showAudioPlayer, setShowAudioPlayer] = useState(true);
+  const isCaptionCurrent = caption.trim() === savedCaption;

  const isItAVideo = isVideo(imageUrl);
  const isItAudio = isAudio(imageUrl);
@ -159,66 +114,49 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({

  return (
    <div className={`flex flex-col ${className}`}>
-      {/* Square image container */}
-      <div
-        ref={cardRef}
-        className="relative w-full"
-        style={{ paddingBottom: '100%' }} // Make it square
-      >
+      <div className="relative w-full" style={{ paddingBottom: '100%' }}>
        <div className="absolute inset-0 rounded-t-lg shadow-md">
-          {inViewport && isVisible && (
-            <>
-              {isItAVideo && (
-                <video
-                  src={`/api/img/${encodeURIComponent(imageUrl)}`}
-                  className={`w-full h-full object-contain`}
-                  autoPlay={false}
-                  loop
-                  muted
-                  controls
-                />
-              )}
-              {isItAudio && !showAudioPlayer && (
-                <div
-                  className="w-full h-full cursor-pointer flex items-center justify-center bg-gray-900"
-                  onClick={() => setShowAudioPlayer(true)}
-                >
-                  <img
-                    src={`/api/audio/art/${encodeURIComponent(imageUrl)}`}
-                    alt={alt}
-                    className="w-full h-full object-contain"
-                    onError={e => {
-                      (e.target as HTMLImageElement).style.display = 'none';
-                    }}
-                  />
-                </div>
-              )}
-              {isItAudio && showAudioPlayer && (
-                <AudioPlayer
-                  src={`/api/img/${encodeURIComponent(imageUrl)}`}
-                  title={imageUrl.replace(/^.*[\\/]/, '')}
-                />
-              )}
-              {isItImage && (
-                <img
-                  src={`/api/img/${encodeURIComponent(imageUrl)}`}
-                  alt={alt}
-                  onLoad={handleLoad}
-                  onClick={onImageClick}
-                  className={classNames('w-full h-full object-contain transition-opacity duration-300', {
-                    'opacity-100': loaded,
-                    'opacity-0': !loaded,
-                    'cursor-zoom-in': !!onImageClick,
-                  })}
-                />
-              )}
-            </>
+          {isItAVideo && (
+            <video
+              src={`/api/img/${encodeURIComponent(imageUrl)}`}
+              className={`w-full h-full object-contain`}
+              autoPlay={false}
+              loop
+              muted
+              controls
+            />
          )}
-          {!isVisible && (
-            <div className="absolute inset-0 flex items-center justify-center bg-gray-800 bg-opacity-75 rounded-t-lg">
-              <span className="text-white text-lg"></span>
+          {isItAudio && !showAudioPlayer && (
+            <div
+              className="w-full h-full cursor-pointer flex items-center justify-center bg-gray-900"
+              onClick={() => setShowAudioPlayer(true)}
+            >
+              <img
+                src={`/api/audio/art/${encodeURIComponent(imageUrl)}`}
+                alt={alt}
+                className="w-full h-full object-contain"
+                onError={e => {
+                  (e.target as HTMLImageElement).style.display = 'none';
+                }}
+              />
            </div>
          )}
+          {isItAudio && showAudioPlayer && (
+            <AudioPlayer src={`/api/img/${encodeURIComponent(imageUrl)}`} title={imageUrl.replace(/^.*[\\/]/, '')} />
+          )}
+          {isItImage && (
+            <img
+              src={`/api/img/${encodeURIComponent(imageUrl)}`}
+              alt={alt}
+              onLoad={handleLoad}
+              onClick={onImageClick}
+              className={classNames('w-full h-full object-contain transition-opacity duration-300', {
+                'opacity-100': loaded,
+                'opacity-0': !loaded,
+                'cursor-zoom-in': !!onImageClick,
+              })}
+            />
+          )}
          {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
@ -254,7 +192,7 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
          'border-transparent border-2': isCaptionCurrent,
        })}
      >
-        {inViewport && isVisible && (isCaptionLoaded || caption) && (
+        {isCaptionLoaded ? (
          <form
            onSubmit={e => {
              e.preventDefault();
@ -269,17 +207,11 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
              value={caption}
              rows={3}
              readOnly={isAutoCaptioning}
-              onChange={e => setCaption(e.target.value)}
+              onChange={e => handleCaptionChange(e.target.value)}
              onKeyDown={handleKeyDown}
            />
          </form>
-        )}
-        {(!inViewport || !isVisible) && isCaptionLoaded && (
-          <div className="w-full h-full flex items-center justify-center text-gray-400">
-            {isVisible ? 'Scroll into view to edit caption' : 'Show content to edit caption'}
-          </div>
-        )}
-        {!isCaptionLoaded && !caption && (
+        ) : (
          <div className="w-full h-full flex items-center justify-center text-gray-400">Loading caption...</div>
        )}
      </div>
--- a/ui/src/components/layout.tsx
+++ b/ui/src/components/layout.tsx
@ -1,4 +1,5 @@
 'use client';
+import React from 'react';
 import classNames from 'classnames';
 import ThemeLogo from './ThemeLogo';
 import { mobileSidebarState } from './Sidebar';
@ -35,10 +36,14 @@ export const TopBar: React.FC<Props> = ({ children, className }) => {
  );
 };

-export const MainContent: React.FC<Props> = ({ children, className }) => {
+export const MainContent = React.forwardRef<HTMLDivElement, Props>(({ children, className }, ref) => {
  return (
-    <div className={classNames('pt-14 px-2 sm:px-4 absolute top-0 left-0 w-full h-full overflow-auto', className)}>
+    <div
+      ref={ref}
+      className={classNames('pt-14 px-2 sm:px-4 absolute top-0 left-0 w-full h-full overflow-auto', className)}
+    >
      {children ? children : null}
    </div>
  );
-};
+});
+MainContent.displayName = 'MainContent';
--- a/ui/src/hooks/useCaptionBatch.tsx
+++ b/ui/src/hooks/useCaptionBatch.tsx
@ -0,0 +1,114 @@
+'use client';
+import { useEffect, useRef, useState } from 'react';
+import { apiClient } from '@/utils/api';
+
+// Module-level batcher: many cards mount at once when the virtualized grid scrolls;
+// instead of N HTTP requests, queue paths and flush them as a single batch.
+type Resolver = { resolve: (caption: string) => void; reject: (err: unknown) => void };
+const pending = new Map<string, Resolver[]>();
+const cache = new Map<string, string>();
+let flushTimer: ReturnType<typeof setTimeout> | null = null;
+const FLUSH_DELAY_MS = 30;
+const MAX_BATCH = 200;
+
+function scheduleFlush() {
+  if (flushTimer) return;
+  flushTimer = setTimeout(flush, FLUSH_DELAY_MS);
+}
+
+async function flush() {
+  flushTimer = null;
+  if (pending.size === 0) return;
+
+  // Drain up to MAX_BATCH paths; if more arrived, reschedule.
+  const paths: string[] = [];
+  for (const path of pending.keys()) {
+    paths.push(path);
+    if (paths.length >= MAX_BATCH) break;
+  }
+  const batchResolvers = paths.map(p => ({ path: p, resolvers: pending.get(p)! }));
+  for (const p of paths) pending.delete(p);
+
+  try {
+    const res = await apiClient.post('/api/caption/getBatch', { imgPaths: paths });
+    const captions: Record<string, string> = res.data?.captions ?? {};
+    for (const { path, resolvers } of batchResolvers) {
+      const value = captions[path] ?? '';
+      cache.set(path, value);
+      for (const r of resolvers) r.resolve(value);
+    }
+  } catch (err) {
+    for (const { resolvers } of batchResolvers) {
+      for (const r of resolvers) r.reject(err);
+    }
+  }
+
+  if (pending.size > 0) scheduleFlush();
+}
+
+function requestCaption(path: string): Promise<string> {
+  return new Promise((resolve, reject) => {
+    const list = pending.get(path);
+    if (list) {
+      list.push({ resolve, reject });
+    } else {
+      pending.set(path, [{ resolve, reject }]);
+    }
+    scheduleFlush();
+  });
+}
+
+export function invalidateCaption(path: string) {
+  cache.delete(path);
+}
+
+export function setCachedCaption(path: string, caption: string) {
+  cache.set(path, caption);
+}
+
+// Fetches caption for a path, using the module-level batcher + cache.
+// `refreshKey` busts the cache (e.g. after external edits or auto-captioning poll).
+export default function useCaptionBatch(imgPath: string | null, refreshKey: number = 0) {
+  const [caption, setCaption] = useState<string>(() => (imgPath ? (cache.get(imgPath) ?? '') : ''));
+  const [isLoaded, setIsLoaded] = useState<boolean>(() => Boolean(imgPath && cache.has(imgPath)));
+  const lastPathRef = useRef<string | null>(null);
+
+  useEffect(() => {
+    if (!imgPath) {
+      setCaption('');
+      setIsLoaded(false);
+      return;
+    }
+
+    if (refreshKey > 0) invalidateCaption(imgPath);
+
+    const cached = cache.get(imgPath);
+    if (cached !== undefined) {
+      setCaption(cached);
+      setIsLoaded(true);
+      lastPathRef.current = imgPath;
+      return;
+    }
+
+    let cancelled = false;
+    lastPathRef.current = imgPath;
+    setIsLoaded(false);
+    requestCaption(imgPath)
+      .then(value => {
+        if (cancelled || lastPathRef.current !== imgPath) return;
+        setCaption(value);
+        setIsLoaded(true);
+      })
+      .catch(err => {
+        if (cancelled) return;
+        console.error('Error fetching caption:', err);
+        setIsLoaded(true);
+      });
+
+    return () => {
+      cancelled = true;
+    };
+  }, [imgPath, refreshKey]);
+
+  return { caption, isLoaded, setCaption };
+}