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 ..

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),

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}`);

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 });
+}

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);
       }
     }

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
       });
 

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));

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 }))}