Added a way to secure the UI. Plus various bug fixes and quality of life updates

ui/src/app/api/auth/route.ts

@@ -0,0 +1,6 @@
+import { NextResponse } from 'next/server';
+
+export async function GET() {
+  // if this gets hit, auth has already been verified
+  return NextResponse.json({ isAuthenticated: true });
+}

ui/src/app/api/datasets/create/route.tsx

@@ -19,4 +19,4 @@ export async function POST(request: Request) {
   } catch (error) {
     return NextResponse.json({ error: 'Failed to create dataset' }, { status: 500 });
   }
-}
+}

ui/src/app/api/datasets/listImages/route.ts

@@ -8,31 +8,25 @@ export async function POST(request: Request) {
   const body = await request.json();
   const { datasetName } = body;
   const datasetFolder = path.join(datasetsPath, datasetName);
-  
+
   try {
     // Check if folder exists
     if (!fs.existsSync(datasetFolder)) {
-      return NextResponse.json(
-        { error: `Folder '${datasetName}' not found` },
-        { status: 404 }
-      );
+      return NextResponse.json({ error: `Folder '${datasetName}' not found` }, { status: 404 });
     }
 
     // Find all images recursively
     const imageFiles = findImagesRecursively(datasetFolder);
-    
+
     // Format response
     const result = imageFiles.map(imgPath => ({
-      img_path: imgPath
+      img_path: imgPath,
     }));
-    
+
     return NextResponse.json({ images: result });
   } catch (error) {
     console.error('Error finding images:', error);
-    return NextResponse.json(
-      { error: 'Failed to process request' },
-      { status: 500 }
-    );
+    return NextResponse.json({ error: 'Failed to process request' }, { status: 500 });
   }
 }
 
@@ -44,13 +38,13 @@ export async function POST(request: Request) {
 function findImagesRecursively(dir: string): string[] {
   const imageExtensions = ['.png', '.jpg', '.jpeg', '.webp'];
   let results: string[] = [];
-  
+
   const items = fs.readdirSync(dir);
-  
+
   for (const item of items) {
     const itemPath = path.join(dir, item);
     const stat = fs.statSync(itemPath);
-    
+
     if (stat.isDirectory()) {
       // If it's a directory, recursively search it
       results = results.concat(findImagesRecursively(itemPath));
@@ -62,6 +56,6 @@ function findImagesRecursively(dir: string): string[] {
       }
     }
   }
-  
+
   return results;
-}
+}

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

@@ -16,7 +16,8 @@ export async function GET(request: NextRequest, { params }: { params: { filePath
     const allowedDirs = [datasetRoot, trainingRoot];
 
     // Security check: Ensure path is in allowed directory
-    const isAllowed = allowedDirs.some(allowedDir => decodedFilePath.startsWith(allowedDir)) && !decodedFilePath.includes('..');
+    const isAllowed =
+      allowedDirs.some(allowedDir => decodedFilePath.startsWith(allowedDir)) && !decodedFilePath.includes('..');
 
     if (!isAllowed) {
       console.warn(`Access denied: ${decodedFilePath} not in ${allowedDirs.join(', ')}`);
@@ -62,7 +63,7 @@ export async function GET(request: NextRequest, { params }: { params: { filePath
       'Accept-Ranges': 'bytes',
       'Cache-Control': 'public, max-age=86400',
       'Content-Disposition': `attachment; filename="${encodeURIComponent(filename)}"`,
-      'X-Content-Type-Options': 'nosniff'
+      'X-Content-Type-Options': 'nosniff',
     };
 
     if (range) {
@@ -70,12 +71,12 @@ export async function GET(request: NextRequest, { params }: { params: { filePath
       const parts = range.replace(/bytes=/, '').split('-');
       const start = parseInt(parts[0], 10);
       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 chunkSize = end - start + 1;
 
-      const fileStream = fs.createReadStream(decodedFilePath, { 
-        start, 
+      const fileStream = fs.createReadStream(decodedFilePath, {
+        start,
         end,
-        highWaterMark: 64 * 1024 // 64KB buffer
+        highWaterMark: 64 * 1024, // 64KB buffer
       });
 
       return new NextResponse(fileStream as any, {
@@ -83,19 +84,19 @@ export async function GET(request: NextRequest, { params }: { params: { filePath
         headers: {
           ...commonHeaders,
           'Content-Range': `bytes ${start}-${end}/${stat.size}`,
-          'Content-Length': String(chunkSize)
+          'Content-Length': String(chunkSize),
         },
       });
     } else {
       // For full file download, read directly without streaming wrapper
       const fileStream = fs.createReadStream(decodedFilePath, {
-        highWaterMark: 64 * 1024 // 64KB buffer
+        highWaterMark: 64 * 1024, // 64KB buffer
       });
 
       return new NextResponse(fileStream as any, {
         headers: {
           ...commonHeaders,
-          'Content-Length': String(stat.size)
+          'Content-Length': String(stat.size),
         },
       });
     }
@@ -103,4 +104,4 @@ export async function GET(request: NextRequest, { params }: { params: { filePath
     console.error('Error serving file:', error);
     return new NextResponse('Internal Server Error', { status: 500 });
   }
-}
+}

ui/src/app/api/gpu/route.ts

@@ -10,7 +10,7 @@ export async function GET() {
     // Get platform
     const platform = os.platform();
     const isWindows = platform === 'win32';
-    
+
     // Check if nvidia-smi is available
     const hasNvidiaSmi = await checkNvidiaSmi(isWindows);
 
@@ -61,8 +61,9 @@ async function checkNvidiaSmi(isWindows: boolean): Promise<boolean> {
 
 async function getGpuStats(isWindows: boolean) {
   // Command is the same for both platforms, but the path might be different
-  const command = 'nvidia-smi --query-gpu=index,name,driver_version,temperature.gpu,utilization.gpu,utilization.memory,memory.total,memory.free,memory.used,power.draw,power.limit,clocks.current.graphics,clocks.current.memory,fan.speed --format=csv,noheader,nounits';
-  
+  const command =
+    'nvidia-smi --query-gpu=index,name,driver_version,temperature.gpu,utilization.gpu,utilization.memory,memory.total,memory.free,memory.used,power.draw,power.limit,clocks.current.graphics,clocks.current.memory,fan.speed --format=csv,noheader,nounits';
+
   // Execute command
   const { stdout } = await execAsync(command);
 
@@ -117,4 +118,4 @@ async function getGpuStats(isWindows: boolean) {
     });
 
   return gpus;
-}
+}

ui/src/app/api/img/caption/route.ts

@@ -17,7 +17,6 @@ export async function POST(request: Request) {
       return NextResponse.json({ error: 'Image does not exist' }, { status: 404 });
     }
 
-
     // check for caption
     const captionPath = imgPath.replace(/\.[^/.]+$/, '') + '.txt';
     // save caption to file

ui/src/app/api/jobs/[jobID]/start/route.ts

@@ -103,7 +103,7 @@ export async function GET(request: NextRequest, { params }: { params: { jobID: s
       cwd: TOOLKIT_ROOT,
       windowsHide: false,
     });
-    
+
     subprocess.unref();
   } else {
     // For non-Windows platforms, use your original approach
@@ -116,7 +116,7 @@ export async function GET(request: NextRequest, { params }: { params: { jobID: s
       },
       cwd: TOOLKIT_ROOT,
     });
-    
+
     subprocess.unref();
   }
   // const subprocess = spawn(pythonPath, [runFilePath, configPath], {

ui/src/app/api/jobs/route.ts

@@ -6,7 +6,6 @@ const prisma = new PrismaClient();
 export async function GET(request: Request) {
   const { searchParams } = new URL(request.url);
   const id = searchParams.get('id');
-  console.log('ID:', id);
 
   try {
     if (id) {
@@ -19,7 +18,6 @@ export async function GET(request: Request) {
     const jobs = await prisma.job.findMany({
       orderBy: { created_at: 'desc' },
     });
-    console.log('Jobs:', jobs);
     return NextResponse.json({ jobs: jobs });
   } catch (error) {
     console.error(error);

ui/src/app/api/settings/route.ts

@@ -1,7 +1,7 @@
 import { NextResponse } from 'next/server';
 import { PrismaClient } from '@prisma/client';
 import { defaultTrainFolder, defaultDatasetsFolder } from '@/paths';
-import {flushCache} from '@/server/settings';
+import { flushCache } from '@/server/settings';
 
 const prisma = new PrismaClient();