github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Jaret Burkett
2026-05-24 15:13:55 -06:00
parent c6a7e81a70
commit 307ff11bc5
12 changed files with 360 additions and 219 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
ui/src/app/api/audio/art/[...audioPath]/route.ts
View File

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

11
ui/src/app/api/caption/get/route.ts
View File

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

45
ui/src/app/api/caption/getBatch/route.ts Normal file
View File

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

24
ui/src/app/api/datasets/listImages/route.ts
View File

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

26
ui/src/app/api/files/[...filePath]/route.ts
View File

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

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

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