Upload 121 files

ui/README.md

@@ -0,0 +1,55 @@
+This is a [Next.js](https://nextjs.org) project bootstrapped with [`create-next-app`](https://nextjs.org/docs/app/api-reference/cli/create-next-app).
+
+## Getting Started
+
+First, run the development server:
+
+```bash
+npm run dev
+# or
+yarn dev
+# or
+pnpm dev
+# or
+bun dev
+```
+
+Open [http://localhost:3000](http://localhost:3000) with your browser to see the result.
+
+You can start editing the page by modifying `app/page.tsx`. The page auto-updates as you edit the file.
+
+This project uses [`next/font`](https://nextjs.org/docs/app/building-your-application/optimizing/fonts) to automatically optimize and load [Geist](https://vercel.com/font), a new font family for Vercel.
+
+## Database Modes
+
+Use the `NEXT_PUBLIC_DB_MODE` environment variable to control how UI data is persisted:
+
+- `server` (default): interacts with the shared SQLite database through Prisma. Supports local job orchestration.
+- `browser`: stores jobs and settings in the user's browser (localStorage). This mode only supports Hugging Face Jobs workflows; local GPU training controls are disabled and the GPU monitor shows a cloud-mode status banner.
+
+When running in browser mode every visitor sees only their own jobs, settings, and dataset catalog (all stored in their browser), making the UI safe to host for multiple users without sharing the SQLite file.
+
+## Hugging Face Authentication
+
+Users can authenticate either by pasting a personal access token or via the Hugging Face OAuth flow. To enable OAuth set the following environment variables for the UI:
+
+- `HF_OAUTH_CLIENT_ID` – the application client ID
+- `HF_OAUTH_CLIENT_SECRET` – the application secret (server-side only)
+- `NEXT_PUBLIC_HF_OAUTH_CLIENT_ID` – the client ID exposed to the browser (usually the same as `HF_OAUTH_CLIENT_ID`)
+
+If these values are not provided the UI falls back to manual token entry. In multi-user/browser mode the authenticated token and namespace are stored per browser session.
+
+## Learn More
+
+To learn more about Next.js, take a look at the following resources:
+
+- [Next.js Documentation](https://nextjs.org/docs) - learn about Next.js features and API.
+- [Learn Next.js](https://nextjs.org/learn) - an interactive Next.js tutorial.
+
+You can check out [the Next.js GitHub repository](https://github.com/vercel/next.js) - your feedback and contributions are welcome!
+
+## Deploy on Vercel
+
+The easiest way to deploy your Next.js app is to use the [Vercel Platform](https://vercel.com/new?utm_medium=default-template&filter=next.js&utm_source=create-next-app&utm_campaign=create-next-app-readme) from the creators of Next.js.
+
+Check out our [Next.js deployment documentation](https://nextjs.org/docs/app/building-your-application/deploying) for more details.

ui/cron/worker.ts

@@ -0,0 +1,31 @@
+class CronWorker {
+  interval: number;
+  is_running: boolean;
+  intervalId: NodeJS.Timeout;
+  constructor() {
+    this.interval = 1000; // Default interval of 1 second
+    this.is_running = false;
+    this.intervalId = setInterval(() => {
+      this.run();
+    }, this.interval);
+  }
+  async run() {
+    if (this.is_running) {
+      return;
+    }
+    this.is_running = true;
+    try {
+      // Loop logic here
+      await this.loop();
+    } catch (error) {
+      console.error('Error in cron worker loop:', error);
+    }
+    this.is_running = false;
+  }
+
+  async loop() {}
+}
+
+// it automatically starts the loop
+const cronWorker = new CronWorker();
+console.log('Cron worker started with interval:', cronWorker.interval, 'ms');

ui/next-env.d.ts

@@ -0,0 +1,5 @@
+/// <reference types="next" />
+/// <reference types="next/image-types/global" />
+
+// NOTE: This file should not be edited
+// see https://nextjs.org/docs/app/api-reference/config/typescript for more information.

ui/next.config.ts

@@ -0,0 +1,15 @@
+import type { NextConfig } from 'next';
+
+const nextConfig: NextConfig = {
+  typescript: {
+    // Remove this. Build fails because of route types
+    ignoreBuildErrors: true,
+  },
+  experimental: {
+    serverActions: {
+      bodySizeLimit: '100mb',
+    },
+  },
+};
+
+export default nextConfig;

ui/package.json

@@ -0,0 +1,51 @@
+{
+  "name": "ai-toolkit-ui",
+  "version": "0.1.0",
+  "private": true,
+  "scripts": {
+    "dev": "concurrently -k -n WORKER,UI \"ts-node-dev --respawn --watch cron --transpile-only cron/worker.ts\" \"next dev --turbopack\"",
+    "build": "tsc -p tsconfig.worker.json && next build",
+    "start": "concurrently --restart-tries -1 --restart-after 1000 -n WORKER,UI \"node dist/worker.js\" \"next start --port 8675\"",
+    "build_and_start": "npm install && npm run update_db && npm run build && npm run start",
+    "lint": "next lint",
+    "update_db": "npx prisma generate && npx prisma db push",
+    "format": "prettier --write \"**/*.{js,jsx,ts,tsx,css,scss}\""
+  },
+  "dependencies": {
+    "@headlessui/react": "^2.2.0",
+    "@huggingface/hub": "^2.5.2",
+    "@monaco-editor/react": "^4.7.0",
+    "@prisma/client": "^6.3.1",
+    "archiver": "^7.0.1",
+    "axios": "^1.7.9",
+    "classnames": "^2.5.1",
+    "form-data": "^4.0.4",
+    "lucide-react": "^0.475.0",
+    "next": "15.1.7",
+    "node-cache": "^5.1.2",
+    "prisma": "^6.3.1",
+    "react": "^19.0.0",
+    "react-dom": "^19.0.0",
+    "react-dropzone": "^14.3.5",
+    "react-global-hooks": "^1.3.5",
+    "react-icons": "^5.5.0",
+    "react-select": "^5.10.1",
+    "sqlite3": "^5.1.7",
+    "uuid": "^11.1.0",
+    "yaml": "^2.7.0"
+  },
+  "devDependencies": {
+    "@types/archiver": "^6.0.3",
+    "@types/node": "^20",
+    "@types/react": "^19",
+    "@types/react-dom": "^19",
+    "concurrently": "^9.1.2",
+    "postcss": "^8",
+    "prettier": "^3.5.1",
+    "prettier-basic": "^1.0.0",
+    "tailwindcss": "^3.4.1",
+    "ts-node-dev": "^2.0.0",
+    "typescript": "^5"
+  },
+  "prettier": "prettier-basic"
+}

ui/postcss.config.mjs

@@ -0,0 +1,8 @@
+/** @type {import('postcss-load-config').Config} */
+const config = {
+  plugins: {
+    tailwindcss: {},
+  },
+};
+
+export default config;

ui/prisma/schema.prisma

@@ -0,0 +1,38 @@
+generator client {
+  provider = "prisma-client-js"
+}
+
+datasource db {
+  provider = "sqlite"
+  url      = "file:../../aitk_db.db"
+}
+
+model Settings {
+  id    Int    @id @default(autoincrement())
+  key   String @unique
+  value String
+}
+
+model Job {
+  id         String   @id @default(uuid())
+  name       String   @unique
+  gpu_ids    String
+  job_config String // JSON string
+  created_at DateTime @default(now())
+  updated_at DateTime @updatedAt
+  status     String   @default("stopped")
+  stop       Boolean  @default(false)
+  step       Int      @default(0)
+  info       String   @default("")
+  speed_string String @default("")
+}
+
+model Queue {
+  id         String   @id @default(uuid())
+  channel    String
+  job_id     String
+  created_at DateTime @default(now())
+  updated_at DateTime @updatedAt
+  status     String   @default("waiting")
+  @@index([job_id, channel])
+}

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

@@ -0,0 +1,112 @@
+import { NextRequest, NextResponse } from 'next/server';
+import { cookies } from 'next/headers';
+
+const TOKEN_ENDPOINT = 'https://huggingface.co/oauth/token';
+const USERINFO_ENDPOINT = 'https://huggingface.co/oauth/userinfo';
+const STATE_COOKIE = 'hf_oauth_state';
+
+function htmlResponse(script: string) {
+  return new NextResponse(
+    `<!DOCTYPE html><html><body><script>${script}</script></body></html>`,
+    {
+      headers: { 'Content-Type': 'text/html; charset=utf-8' },
+    },
+  );
+}
+
+export async function GET(request: NextRequest) {
+  const clientId = process.env.HF_OAUTH_CLIENT_ID || process.env.NEXT_PUBLIC_HF_OAUTH_CLIENT_ID;
+  const clientSecret = process.env.HF_OAUTH_CLIENT_SECRET;
+
+  if (!clientId || !clientSecret) {
+    return NextResponse.json({ error: 'OAuth application is not configured' }, { status: 500 });
+  }
+
+  const { searchParams } = new URL(request.url);
+  const code = searchParams.get('code');
+  const incomingState = searchParams.get('state');
+
+  const cookieStore = cookies();
+  const storedState = cookieStore.get(STATE_COOKIE)?.value;
+
+  cookieStore.delete(STATE_COOKIE);
+
+  const origin = request.nextUrl.origin;
+
+  if (!code || !incomingState || !storedState || incomingState !== storedState) {
+    const script = `
+      window.opener && window.opener.postMessage({
+        type: 'HF_OAUTH_ERROR',
+        payload: { message: 'Invalid or expired OAuth state.' }
+      }, '${origin}');
+      window.close();
+    `;
+    return htmlResponse(script.trim());
+  }
+
+  const redirectUri = `${origin}/api/auth/hf/callback`;
+
+  try {
+    const tokenResponse = await fetch(TOKEN_ENDPOINT, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/x-www-form-urlencoded',
+      },
+      body: new URLSearchParams({
+        grant_type: 'authorization_code',
+        code,
+        redirect_uri: redirectUri,
+        client_id: clientId,
+        client_secret: clientSecret,
+      }),
+    });
+
+    if (!tokenResponse.ok) {
+      const errorPayload = await tokenResponse.json().catch(() => ({}));
+      throw new Error(errorPayload?.error_description || 'Failed to exchange code for token');
+    }
+
+    const tokenData = await tokenResponse.json();
+    const accessToken = tokenData?.access_token;
+    if (!accessToken) {
+      throw new Error('Access token missing in response');
+    }
+
+    const userResponse = await fetch(USERINFO_ENDPOINT, {
+      headers: {
+        Authorization: `Bearer ${accessToken}`,
+      },
+    });
+
+    if (!userResponse.ok) {
+      throw new Error('Failed to fetch user info');
+    }
+
+    const profile = await userResponse.json();
+    const namespace = profile?.preferred_username || profile?.name || 'user';
+
+    const script = `
+      window.opener && window.opener.postMessage({
+        type: 'HF_OAUTH_SUCCESS',
+        payload: {
+          token: ${JSON.stringify(accessToken)},
+          namespace: ${JSON.stringify(namespace)},
+        }
+      }, '${origin}');
+      window.close();
+    `;
+
+    return htmlResponse(script.trim());
+  } catch (error: any) {
+    const message = error?.message || 'OAuth flow failed';
+    const script = `
+      window.opener && window.opener.postMessage({
+        type: 'HF_OAUTH_ERROR',
+        payload: { message: ${JSON.stringify(message)} }
+      }, '${origin}');
+      window.close();
+    `;
+
+    return htmlResponse(script.trim());
+  }
+}

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

@@ -0,0 +1,36 @@
+import { randomUUID } from 'crypto';
+import { NextRequest, NextResponse } from 'next/server';
+
+const HF_AUTHORIZE_URL = 'https://huggingface.co/oauth/authorize';
+const STATE_COOKIE = 'hf_oauth_state';
+
+export async function GET(request: NextRequest) {
+  const clientId = process.env.HF_OAUTH_CLIENT_ID || process.env.NEXT_PUBLIC_HF_OAUTH_CLIENT_ID;
+  if (!clientId) {
+    return NextResponse.json({ error: 'OAuth client ID not configured' }, { status: 500 });
+  }
+
+  const state = randomUUID();
+  const origin = request.nextUrl.origin;
+  const redirectUri = `${origin}/api/auth/hf/callback`;
+
+  const authorizeUrl = new URL(HF_AUTHORIZE_URL);
+  authorizeUrl.searchParams.set('response_type', 'code');
+  authorizeUrl.searchParams.set('client_id', clientId);
+  authorizeUrl.searchParams.set('redirect_uri', redirectUri);
+  authorizeUrl.searchParams.set('scope', 'openid profile read-repos');
+  authorizeUrl.searchParams.set('state', state);
+
+  const response = NextResponse.redirect(authorizeUrl.toString(), { status: 302 });
+  response.cookies.set({
+    name: STATE_COOKIE,
+    value: state,
+    httpOnly: true,
+    sameSite: 'lax',
+    secure: process.env.NODE_ENV === 'production',
+    maxAge: 60 * 5,
+    path: '/',
+  });
+
+  return response;
+}

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

@@ -0,0 +1,22 @@
+import { NextRequest, NextResponse } from 'next/server';
+import { whoAmI } from '@huggingface/hub';
+
+export async function POST(request: NextRequest) {
+  try {
+    const body = await request.json().catch(() => ({}));
+    const token = (body?.token || '').trim();
+
+    if (!token) {
+      return NextResponse.json({ error: 'Token is required' }, { status: 400 });
+    }
+
+    const info = await whoAmI({ accessToken: token });
+    return NextResponse.json({
+      name: info?.name || info?.username || 'user',
+      email: info?.email || null,
+      orgs: info?.orgs || [],
+    });
+  } catch (error: any) {
+    return NextResponse.json({ error: error?.message || 'Invalid token' }, { status: 401 });
+  }
+}

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/caption/get/route.ts

@@ -0,0 +1,46 @@
+/* eslint-disable */
+import { NextRequest, NextResponse } from 'next/server';
+import fs from 'fs';
+import path from 'path';
+import { getDatasetsRoot } from '@/server/settings';
+
+export async function POST(request: NextRequest) {
+  
+  const body = await request.json();
+  const { imgPath } = body;
+  console.log('Received POST request for caption:', imgPath);
+  try {
+    // Decode the path
+    const filepath = imgPath;
+    console.log('Decoded image path:', filepath);
+
+    // caption name is the filepath without extension but with .txt
+    const captionPath = filepath.replace(/\.[^/.]+$/, '') + '.txt';
+
+    // Get allowed directories
+    const allowedDir = await getDatasetsRoot();
+
+    // Security check: Ensure path is in allowed directory
+    const isAllowed = filepath.startsWith(allowedDir) && !filepath.includes('..');
+
+    if (!isAllowed) {
+      console.warn(`Access denied: ${filepath} not in ${allowedDir}`);
+      return new NextResponse('Access denied', { status: 403 });
+    }
+
+    // Check if file exists
+    if (!fs.existsSync(captionPath)) {
+      // send back blank string if caption file does not exist
+      return new NextResponse('');
+    }
+
+    // Read caption file
+    const caption = fs.readFileSync(captionPath, 'utf-8');
+
+    // Return caption
+    return new NextResponse(caption);
+  } catch (error) {
+    console.error('Error getting caption:', error);
+    return new NextResponse('Error getting caption', { status: 500 });
+  }
+}

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

@@ -0,0 +1,25 @@
+import { NextResponse } from 'next/server';
+import fs from 'fs';
+import path from 'path';
+import { getDatasetsRoot } from '@/server/settings';
+
+export async function POST(request: Request) {
+  try {
+    const body = await request.json();
+    let { name } = body;
+    // clean name by making lower case,  removing special characters, and replacing spaces with underscores
+    name = name.toLowerCase().replace(/[^a-z0-9]+/g, '_');
+
+    let datasetsPath = await getDatasetsRoot();
+    let datasetPath = path.join(datasetsPath, name);
+
+    // if folder doesnt exist, create it
+    if (!fs.existsSync(datasetPath)) {
+      fs.mkdirSync(datasetPath);
+    }
+
+    return NextResponse.json({ success: true, name: name, path: datasetPath });
+  } catch (error) {
+    return NextResponse.json({ error: 'Failed to create dataset' }, { status: 500 });
+  }
+}

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

@@ -0,0 +1,24 @@
+import { NextResponse } from 'next/server';
+import fs from 'fs';
+import path from 'path';
+import { getDatasetsRoot } from '@/server/settings';
+
+export async function POST(request: Request) {
+  try {
+    const body = await request.json();
+    const { name } = body;
+    let datasetsPath = await getDatasetsRoot();
+    let datasetPath = path.join(datasetsPath, name);
+
+    // if folder doesnt exist, ignore
+    if (!fs.existsSync(datasetPath)) {
+      return NextResponse.json({ success: true });
+    }
+
+    // delete it and return success
+    fs.rmdirSync(datasetPath, { recursive: true });
+    return NextResponse.json({ success: true });
+  } catch (error) {
+    return NextResponse.json({ error: 'Failed to create dataset' }, { status: 500 });
+  }
+}

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

@@ -0,0 +1,25 @@
+import { NextResponse } from 'next/server';
+import fs from 'fs';
+import { getDatasetsRoot } from '@/server/settings';
+
+export async function GET() {
+  try {
+    let datasetsPath = await getDatasetsRoot();
+
+    // if folder doesnt exist, create it
+    if (!fs.existsSync(datasetsPath)) {
+      fs.mkdirSync(datasetsPath);
+    }
+
+    // find all the folders in the datasets folder
+    let folders = fs
+      .readdirSync(datasetsPath, { withFileTypes: true })
+      .filter(dirent => dirent.isDirectory())
+      .filter(dirent => !dirent.name.startsWith('.'))
+      .map(dirent => dirent.name);
+
+    return NextResponse.json(folders);
+  } catch (error) {
+    return NextResponse.json({ error: 'Failed to fetch datasets' }, { status: 500 });
+  }
+}

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

@@ -0,0 +1,61 @@
+import { NextResponse } from 'next/server';
+import fs from 'fs';
+import path from 'path';
+import { getDatasetsRoot } from '@/server/settings';
+
+export async function POST(request: Request) {
+  const datasetsPath = await getDatasetsRoot();
+  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 });
+    }
+
+    // Find all images recursively
+    const imageFiles = findImagesRecursively(datasetFolder);
+
+    // Format response
+    const result = imageFiles.map(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 });
+  }
+}
+
+/**
+ * Recursively finds all image files in a directory and its subdirectories
+ * @param dir Directory to search
+ * @returns Array of absolute paths to image files
+ */
+function findImagesRecursively(dir: string): string[] {
+  const imageExtensions = ['.png', '.jpg', '.jpeg', '.webp', '.mp4', '.avi', '.mov', '.mkv', '.wmv', '.m4v', '.flv'];
+  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() && item !== '_controls' && !item.startsWith('.')) {
+      // If it's a directory, recursively search it
+      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)) {
+        results.push(itemPath);
+      }
+    }
+  }
+
+  return results;
+}

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

@@ -0,0 +1,57 @@
+// src/app/api/datasets/upload/route.ts
+import { NextRequest, NextResponse } from 'next/server';
+import { writeFile, mkdir } from 'fs/promises';
+import { join } from 'path';
+import { getDatasetsRoot } from '@/server/settings';
+
+export async function POST(request: NextRequest) {
+  try {
+    const datasetsPath = await getDatasetsRoot();
+    if (!datasetsPath) {
+      return NextResponse.json({ error: 'Datasets path not found' }, { status: 500 });
+    }
+    const formData = await request.formData();
+    const files = formData.getAll('files');
+    const datasetName = formData.get('datasetName') as string;
+
+    if (!files || files.length === 0) {
+      return NextResponse.json({ error: 'No files provided' }, { status: 400 });
+    }
+
+    // Create upload directory if it doesn't exist
+    const uploadDir = join(datasetsPath, datasetName);
+    await mkdir(uploadDir, { recursive: true });
+
+    const savedFiles: string[] = [];
+    
+    // Process files sequentially to avoid overwhelming the system
+    for (let i = 0; i < files.length; i++) {
+      const file = files[i] as any;
+      const bytes = await file.arrayBuffer();
+      const buffer = Buffer.from(bytes);
+
+      // Clean filename and ensure it's unique
+      const fileName = file.name.replace(/[^a-zA-Z0-9.-]/g, '_');
+      const filePath = join(uploadDir, fileName);
+
+      await writeFile(filePath, buffer);
+      savedFiles.push(fileName);
+    }
+
+    return NextResponse.json({
+      message: 'Files uploaded successfully',
+      files: savedFiles,
+    });
+  } catch (error) {
+    console.error('Upload error:', error);
+    return NextResponse.json({ error: 'Error uploading files' }, { status: 500 });
+  }
+}
+
+// Increase payload size limit (default is 4mb)
+export const config = {
+  api: {
+    bodyParser: false,
+    responseLimit: '50mb',
+  },
+};

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

@@ -0,0 +1,116 @@
+/* eslint-disable */
+import { NextRequest, NextResponse } from 'next/server';
+import fs from 'fs';
+import path from 'path';
+import { getDatasetsRoot, getTrainingFolder } from '@/server/settings';
+
+export async function GET(request: NextRequest, { params }: { params: { filePath: string } }) {
+  const { filePath } = await params;
+  try {
+    // Decode the path
+    const decodedFilePath = decodeURIComponent(filePath);
+
+    // Get allowed directories
+    const datasetRoot = await getDatasetsRoot();
+    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('..');
+
+    if (!isAllowed) {
+      console.warn(`Access denied: ${decodedFilePath} 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}`);
+      return new NextResponse('File not found', { status: 404 });
+    }
+
+    // Get file info
+    const stat = fs.statSync(decodedFilePath);
+    if (!stat.isFile()) {
+      return new NextResponse('Not a file', { status: 400 });
+    }
+
+    // Get filename for Content-Disposition
+    const filename = path.basename(decodedFilePath);
+
+    // Determine content type
+    const ext = path.extname(decodedFilePath).toLowerCase();
+    const contentTypeMap: { [key: string]: string } = {
+      '.jpg': 'image/jpeg',
+      '.jpeg': 'image/jpeg',
+      '.png': 'image/png',
+      '.gif': 'image/gif',
+      '.webp': 'image/webp',
+      '.svg': 'image/svg+xml',
+      '.bmp': 'image/bmp',
+      '.safetensors': 'application/octet-stream',
+      '.zip': 'application/zip',
+      // Videos
+      '.mp4': 'video/mp4',
+      '.avi': 'video/x-msvideo',
+      '.mov': 'video/quicktime',
+      '.mkv': 'video/x-matroska',
+      '.wmv': 'video/x-ms-wmv',
+      '.m4v': 'video/x-m4v',
+      '.flv': 'video/x-flv'
+    };
+
+    const contentType = contentTypeMap[ext] || 'application/octet-stream';
+
+    // Get range header for partial content support
+    const range = request.headers.get('range');
+
+    // Common headers for better download handling
+    const commonHeaders = {
+      'Content-Type': contentType,
+      'Accept-Ranges': 'bytes',
+      'Cache-Control': 'public, max-age=86400',
+      'Content-Disposition': `attachment; filename="${encodeURIComponent(filename)}"`,
+      'X-Content-Type-Options': 'nosniff',
+    };
+
+    if (range) {
+      // Parse range header
+      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 fileStream = fs.createReadStream(decodedFilePath, {
+        start,
+        end,
+        highWaterMark: 64 * 1024, // 64KB buffer
+      });
+
+      return new NextResponse(fileStream as any, {
+        status: 206,
+        headers: {
+          ...commonHeaders,
+          'Content-Range': `bytes ${start}-${end}/${stat.size}`,
+          'Content-Length': String(chunkSize),
+        },
+      });
+    } else {
+      // For full file download, read directly without streaming wrapper
+      const fileStream = fs.createReadStream(decodedFilePath, {
+        highWaterMark: 64 * 1024, // 64KB buffer
+      });
+
+      return new NextResponse(fileStream as any, {
+        headers: {
+          ...commonHeaders,
+          'Content-Length': String(stat.size),
+        },
+      });
+    }
+  } catch (error) {
+    console.error('Error serving file:', error);
+    return new NextResponse('Internal Server Error', { status: 500 });
+  }
+}

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

@@ -0,0 +1,121 @@
+import { NextResponse } from 'next/server';
+import { exec } from 'child_process';
+import { promisify } from 'util';
+import os from 'os';
+
+const execAsync = promisify(exec);
+
+export async function GET() {
+  try {
+    // Get platform
+    const platform = os.platform();
+    const isWindows = platform === 'win32';
+
+    // Check if nvidia-smi is available
+    const hasNvidiaSmi = await checkNvidiaSmi(isWindows);
+
+    if (!hasNvidiaSmi) {
+      return NextResponse.json({
+        hasNvidiaSmi: false,
+        gpus: [],
+        error: 'nvidia-smi not found or not accessible',
+      });
+    }
+
+    // Get GPU stats
+    const gpuStats = await getGpuStats(isWindows);
+
+    return NextResponse.json({
+      hasNvidiaSmi: true,
+      gpus: gpuStats,
+    });
+  } catch (error) {
+    console.error('Error fetching NVIDIA GPU stats:', error);
+    return NextResponse.json(
+      {
+        hasNvidiaSmi: false,
+        gpus: [],
+        error: `Failed to fetch GPU stats: ${error instanceof Error ? error.message : String(error)}`,
+      },
+      { status: 500 },
+    );
+  }
+}
+
+async function checkNvidiaSmi(isWindows: boolean): Promise<boolean> {
+  try {
+    if (isWindows) {
+      // Check if nvidia-smi is available on Windows
+      // It's typically located in C:\Program Files\NVIDIA Corporation\NVSMI\nvidia-smi.exe
+      // but we'll just try to run it directly as it may be in PATH
+      await execAsync('nvidia-smi -L');
+    } else {
+      // Linux/macOS check
+      await execAsync('which nvidia-smi');
+    }
+    return true;
+  } catch (error) {
+    return false;
+  }
+}
+
+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';
+
+  // Execute command
+  const { stdout } = await execAsync(command);
+
+  // Parse CSV output
+  const gpus = stdout
+    .trim()
+    .split('\n')
+    .map(line => {
+      const [
+        index,
+        name,
+        driverVersion,
+        temperature,
+        gpuUtil,
+        memoryUtil,
+        memoryTotal,
+        memoryFree,
+        memoryUsed,
+        powerDraw,
+        powerLimit,
+        clockGraphics,
+        clockMemory,
+        fanSpeed,
+      ] = line.split(', ').map(item => item.trim());
+
+      return {
+        index: parseInt(index),
+        name,
+        driverVersion,
+        temperature: parseInt(temperature),
+        utilization: {
+          gpu: parseInt(gpuUtil),
+          memory: parseInt(memoryUtil),
+        },
+        memory: {
+          total: parseInt(memoryTotal),
+          free: parseInt(memoryFree),
+          used: parseInt(memoryUsed),
+        },
+        power: {
+          draw: parseFloat(powerDraw),
+          limit: parseFloat(powerLimit),
+        },
+        clocks: {
+          graphics: parseInt(clockGraphics),
+          memory: parseInt(clockMemory),
+        },
+        fan: {
+          speed: parseInt(fanSpeed) || 0, // Some GPUs might not report fan speed, default to 0
+        },
+      };
+    });
+
+  return gpus;
+}

ui/src/app/api/hf-hub/route.ts

@@ -0,0 +1,165 @@
+import { NextRequest, NextResponse } from 'next/server';
+import { whoAmI, createRepo, uploadFiles, datasetInfo } from '@huggingface/hub';
+import { readdir, stat } from 'fs/promises';
+import path from 'path';
+
+export async function POST(request: NextRequest) {
+  try {
+    const body = await request.json();
+    const { action, token, namespace, datasetName, datasetPath, datasetId } = body;
+
+    if (!token) {
+      return NextResponse.json({ error: 'HF token is required' }, { status: 400 });
+    }
+
+    switch (action) {
+      case 'whoami':
+        try {
+          const user = await whoAmI({ accessToken: token });
+          return NextResponse.json({ user });
+        } catch (error) {
+          return NextResponse.json({ error: 'Invalid token or network error' }, { status: 401 });
+        }
+
+      case 'createDataset':
+        try {
+          if (!namespace || !datasetName) {
+            return NextResponse.json({ error: 'Namespace and dataset name required' }, { status: 400 });
+          }
+
+          const repoId = `datasets/${namespace}/${datasetName}`;
+          
+          // Create repository
+          await createRepo({
+            repo: repoId,
+            accessToken: token,
+            private: false,
+          });
+
+          return NextResponse.json({ success: true, repoId });
+        } catch (error: any) {
+          if (error.message?.includes('already exists')) {
+            return NextResponse.json({ success: true, repoId: `${namespace}/${datasetName}`, exists: true });
+          }
+          return NextResponse.json({ error: error.message || 'Failed to create dataset' }, { status: 500 });
+        }
+
+      case 'uploadDataset':
+        try {
+          if (!namespace || !datasetName || !datasetPath) {
+            return NextResponse.json({ error: 'Missing required parameters' }, { status: 400 });
+          }
+
+          const repoId = `datasets/${namespace}/${datasetName}`;
+          
+          // Check if directory exists
+          try {
+            await stat(datasetPath);
+          } catch {
+            return NextResponse.json({ error: 'Dataset path does not exist' }, { status: 400 });
+          }
+
+          // Read files from directory and upload them
+          const files = await readdir(datasetPath);
+          const filesToUpload = [];
+          
+          for (const fileName of files) {
+            const filePath = path.join(datasetPath, fileName);
+            const fileStats = await stat(filePath);
+            
+            if (fileStats.isFile()) {
+              filesToUpload.push({
+                path: fileName,
+                content: new URL(`file://${filePath}`)
+              });
+            }
+          }
+
+          if (filesToUpload.length > 0) {
+            await uploadFiles({
+              repo: repoId,
+              accessToken: token,
+              files: filesToUpload,
+            });
+          }
+
+          return NextResponse.json({ success: true, repoId });
+        } catch (error: any) {
+          console.error('Upload error:', error);
+          return NextResponse.json({ error: error.message || 'Failed to upload dataset' }, { status: 500 });
+        }
+
+      case 'listFiles':
+        try {
+          if (!datasetPath) {
+            return NextResponse.json({ error: 'Dataset path required' }, { status: 400 });
+          }
+
+          const files = await readdir(datasetPath, { withFileTypes: true });
+          const imageExtensions = ['.jpg', '.jpeg', '.png', '.webp', '.bmp'];
+          
+          const imageFiles = files
+            .filter(file => file.isFile())
+            .filter(file => imageExtensions.some(ext => file.name.toLowerCase().endsWith(ext)))
+            .map(file => ({
+              name: file.name,
+              path: path.join(datasetPath, file.name),
+            }));
+
+          const captionFiles = files
+            .filter(file => file.isFile())
+            .filter(file => file.name.endsWith('.txt'))
+            .map(file => ({
+              name: file.name,
+              path: path.join(datasetPath, file.name),
+            }));
+
+          return NextResponse.json({ 
+            images: imageFiles, 
+            captions: captionFiles,
+            total: imageFiles.length 
+          });
+        } catch (error: any) {
+          return NextResponse.json({ error: error.message || 'Failed to list files' }, { status: 500 });
+        }
+
+      case 'validateDataset':
+        try {
+          if (!datasetId) {
+            return NextResponse.json({ error: 'Dataset ID required' }, { status: 400 });
+          }
+
+          // Try to get dataset info to validate it exists and is accessible
+          const dataset = await datasetInfo({
+            name: datasetId,
+            accessToken: token,
+          });
+
+          return NextResponse.json({ 
+            exists: true, 
+            dataset: {
+              id: dataset.id,
+              author: dataset.author,
+              downloads: dataset.downloads,
+              likes: dataset.likes,
+              private: dataset.private,
+            }
+          });
+        } catch (error: any) {
+          if (error.message?.includes('404') || error.message?.includes('not found')) {
+            return NextResponse.json({ exists: false }, { status: 200 });
+          }
+          if (error.message?.includes('401') || error.message?.includes('403')) {
+            return NextResponse.json({ error: 'Dataset not accessible with current token' }, { status: 403 });
+          }
+          return NextResponse.json({ error: error.message || 'Failed to validate dataset' }, { status: 500 });
+        }
+
+      default:
+        return NextResponse.json({ error: 'Invalid action' }, { status: 400 });
+    }
+  } catch (error: any) {
+    console.error('HF Hub API error:', error);
+    return NextResponse.json({ error: error.message || 'Internal server error' }, { status: 500 });
+  }
+}

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

@@ -0,0 +1,761 @@
+import { NextRequest, NextResponse } from 'next/server';
+import { spawn } from 'child_process';
+import { writeFile } from 'fs/promises';
+import path from 'path';
+import { tmpdir } from 'os';
+
+export async function POST(request: NextRequest) {
+  try {
+    const body = await request.json();
+    const { action, token, hardware, namespace, jobConfig, datasetRepo } = body;
+
+    switch (action) {
+      case 'checkStatus':
+        try {
+          if (!token || !jobConfig?.hf_job_id) {
+            return NextResponse.json({ error: 'Token and job ID required' }, { status: 400 });
+          }
+
+          const jobStatus = await checkHFJobStatus(token, jobConfig.hf_job_id);
+          return NextResponse.json({ status: jobStatus });
+        } catch (error: any) {
+          console.error('Job status check error:', error);
+          return NextResponse.json({ error: error.message }, { status: 500 });
+        }
+
+      case 'generateScript':
+        try {
+          const uvScript = generateUVScript({
+            jobConfig,
+            datasetRepo,
+            namespace,
+            token: token || 'YOUR_HF_TOKEN',
+          });
+
+          return NextResponse.json({ 
+            script: uvScript,
+            filename: `train_${jobConfig.config.name.replace(/[^a-zA-Z0-9]/g, '_')}.py`
+          });
+        } catch (error: any) {
+          return NextResponse.json({ error: error.message }, { status: 500 });
+        }
+
+      case 'submitJob':
+        try {
+          if (!token || !hardware) {
+            return NextResponse.json({ error: 'Token and hardware required' }, { status: 400 });
+          }
+
+          // Generate UV script
+          const uvScript = generateUVScript({
+            jobConfig,
+            datasetRepo,
+            namespace,
+            token,
+          });
+
+          // Write script to temporary file
+          const scriptPath = path.join(tmpdir(), `train_${Date.now()}.py`);
+          await writeFile(scriptPath, uvScript);
+
+          // Submit HF job using uv run
+          const jobId = await submitHFJobUV(token, hardware, scriptPath);
+
+          return NextResponse.json({ 
+            success: true, 
+            jobId,
+            message: `Job submitted successfully with ID: ${jobId}`
+          });
+        } catch (error: any) {
+          console.error('Job submission error:', error);
+          return NextResponse.json({ error: error.message }, { status: 500 });
+        }
+
+      default:
+        return NextResponse.json({ error: 'Invalid action' }, { status: 400 });
+    }
+  } catch (error: any) {
+    console.error('HF Jobs API error:', error);
+    return NextResponse.json({ error: error.message }, { status: 500 });
+  }
+}
+
+function generateUVScript({ jobConfig, datasetRepo, namespace, token }: {
+  jobConfig: any;
+  datasetRepo: string;
+  namespace: string;
+  token: string;
+}) {
+  const config = jobConfig.config;
+  const process = config.process[0];
+
+  return `# /// script
+# dependencies = [
+#     "torch>=2.0.0",
+#     "torchvision",
+#     "torchao==0.10.0",
+#     "safetensors",
+#     "diffusers @ git+https://github.com/huggingface/diffusers@7a2b78bf0f788d311cc96b61e660a8e13e3b1e63",
+#     "transformers==4.52.4",
+#     "lycoris-lora==1.8.3",
+#     "flatten_json",
+#     "pyyaml",
+#     "oyaml",
+#     "tensorboard",
+#     "kornia",
+#     "invisible-watermark",
+#     "einops",
+#     "accelerate",
+#     "toml",
+#     "albumentations==1.4.15",
+#     "albucore==0.0.16",
+#     "pydantic",
+#     "omegaconf",
+#     "k-diffusion",
+#     "open_clip_torch",
+#     "timm",
+#     "prodigyopt",
+#     "controlnet_aux==0.0.10",
+#     "python-dotenv",
+#     "bitsandbytes",
+#     "hf_transfer",
+#     "lpips",
+#     "pytorch_fid",
+#     "optimum-quanto==0.2.4",
+#     "sentencepiece",
+#     "huggingface_hub",
+#     "peft",
+#     "python-slugify",
+#     "opencv-python-headless",
+#     "pytorch-wavelets==1.3.0",
+#     "matplotlib==3.10.1",
+#     "setuptools==69.5.1",
+#     "datasets==4.0.0",
+#     "pyarrow==20.0.0",
+#     "pillow",
+#     "ftfy",
+# ]
+# ///
+
+import os
+import sys
+import subprocess
+import argparse
+import oyaml as yaml
+from datasets import load_dataset
+from huggingface_hub import HfApi, create_repo, upload_folder, snapshot_download
+import tempfile
+import shutil
+import glob
+from PIL import Image
+
+def setup_ai_toolkit():
+    """Clone and setup ai-toolkit repository"""
+    repo_dir = "ai-toolkit"
+    if not os.path.exists(repo_dir):
+        print("Cloning ai-toolkit repository...")
+        subprocess.run(
+            ["git", "clone", "https://github.com/ostris/ai-toolkit.git", repo_dir],
+            check=True
+        )
+    sys.path.insert(0, os.path.abspath(repo_dir))
+    return repo_dir
+
+def download_dataset(dataset_repo: str, local_path: str):
+    """Download dataset from HF Hub as files"""
+    print(f"Downloading dataset from {dataset_repo}...")
+    
+    # Create local dataset directory
+    os.makedirs(local_path, exist_ok=True)
+    
+    # Use snapshot_download to get the dataset files directly
+    from huggingface_hub import snapshot_download
+    
+    try:
+        # First try to download as a structured dataset
+        dataset = load_dataset(dataset_repo, split="train")
+        
+        # Download images and captions from structured dataset
+        for i, item in enumerate(dataset):
+            # Save image
+            if "image" in item:
+                image_path = os.path.join(local_path, f"image_{i:06d}.jpg")
+                image = item["image"]
+                
+                # Convert RGBA to RGB if necessary (for JPEG compatibility)
+                if image.mode == 'RGBA':
+                    # Create a white background and paste the RGBA image on it
+                    background = Image.new('RGB', image.size, (255, 255, 255))
+                    background.paste(image, mask=image.split()[-1])  # Use alpha channel as mask
+                    image = background
+                elif image.mode not in ['RGB', 'L']:
+                    # Convert any other mode to RGB
+                    image = image.convert('RGB')
+                
+                image.save(image_path, 'JPEG')
+            
+            # Save caption
+            if "text" in item:
+                caption_path = os.path.join(local_path, f"image_{i:06d}.txt")
+                with open(caption_path, "w", encoding="utf-8") as f:
+                    f.write(item["text"])
+        
+        print(f"Downloaded {len(dataset)} items to {local_path}")
+        
+    except Exception as e:
+        print(f"Failed to load as structured dataset: {e}")
+        print("Attempting to download raw files...")
+        
+        # Download the dataset repository as files
+        temp_repo_path = snapshot_download(repo_id=dataset_repo, repo_type="dataset")
+        
+        # Copy all image and text files to the local path
+        import glob
+        import shutil
+        
+        print(f"Downloaded repo to: {temp_repo_path}")
+        print(f"Contents: {os.listdir(temp_repo_path)}")
+        
+        # Find all image files
+        image_extensions = ['*.jpg', '*.jpeg', '*.png', '*.webp', '*.bmp', '*.JPG', '*.JPEG', '*.PNG']
+        image_files = []
+        for ext in image_extensions:
+            pattern = os.path.join(temp_repo_path, "**", ext)
+            found_files = glob.glob(pattern, recursive=True)
+            image_files.extend(found_files)
+            print(f"Pattern {pattern} found {len(found_files)} files")
+        
+        # Find all text files
+        text_files = glob.glob(os.path.join(temp_repo_path, "**", "*.txt"), recursive=True)
+        
+        print(f"Found {len(image_files)} image files and {len(text_files)} text files")
+        
+        # Copy image files
+        for i, img_file in enumerate(image_files):
+            dest_path = os.path.join(local_path, f"image_{i:06d}.jpg")
+            
+            # Load and convert image if needed
+            try:
+                with Image.open(img_file) as image:
+                    if image.mode == 'RGBA':
+                        background = Image.new('RGB', image.size, (255, 255, 255))
+                        background.paste(image, mask=image.split()[-1])
+                        image = background
+                    elif image.mode not in ['RGB', 'L']:
+                        image = image.convert('RGB')
+                    
+                    image.save(dest_path, 'JPEG')
+            except Exception as img_error:
+                print(f"Error processing image {img_file}: {img_error}")
+                continue
+        
+        # Copy text files (captions)
+        for i, txt_file in enumerate(text_files[:len(image_files)]):  # Match number of images
+            dest_path = os.path.join(local_path, f"image_{i:06d}.txt")
+            try:
+                shutil.copy2(txt_file, dest_path)
+            except Exception as txt_error:
+                print(f"Error copying text file {txt_file}: {txt_error}")
+                continue
+        
+        print(f"Downloaded {len(image_files)} images and {len(text_files)} captions to {local_path}")
+
+def create_config(dataset_path: str, output_path: str):
+    """Create training configuration"""
+    import json
+    
+    # Load config from JSON string and fix boolean/null values for Python
+    config_str = """${JSON.stringify(jobConfig, null, 2)}"""
+    config_str = config_str.replace('true', 'True').replace('false', 'False').replace('null', 'None')
+    config = eval(config_str)
+    
+    # Update paths for cloud environment
+    config["config"]["process"][0]["datasets"][0]["folder_path"] = dataset_path
+    config["config"]["process"][0]["training_folder"] = output_path
+    
+    # Remove sqlite_db_path as it's not needed for cloud training
+    if "sqlite_db_path" in config["config"]["process"][0]:
+        del config["config"]["process"][0]["sqlite_db_path"]
+    
+    # Also change trainer type from ui_trainer to standard trainer to avoid UI dependencies
+    if config["config"]["process"][0]["type"] == "ui_trainer":
+        config["config"]["process"][0]["type"] = "sd_trainer"
+    
+    return config
+
+def upload_results(output_path: str, model_name: str, namespace: str, token: str, config: dict):
+    """Upload trained model to HF Hub with README generation and proper file organization"""
+    import tempfile
+    import shutil
+    import glob
+    import re
+    import yaml
+    from datetime import datetime
+    from huggingface_hub import create_repo, upload_file, HfApi
+    
+    try:
+        repo_id = f"{namespace}/{model_name}"
+        
+        # Create repository
+        create_repo(repo_id=repo_id, token=token, exist_ok=True)
+        
+        print(f"Uploading model to {repo_id}...")
+        
+        # Create temporary directory for organized upload
+        with tempfile.TemporaryDirectory() as temp_upload_dir:
+            api = HfApi()
+            
+            # 1. Find and upload model files to root directory
+            safetensors_files = glob.glob(os.path.join(output_path, "**", "*.safetensors"), recursive=True)
+            json_files = glob.glob(os.path.join(output_path, "**", "*.json"), recursive=True)
+            txt_files = glob.glob(os.path.join(output_path, "**", "*.txt"), recursive=True)
+            
+            uploaded_files = []
+            
+            # Upload .safetensors files to root
+            for file_path in safetensors_files:
+                filename = os.path.basename(file_path)
+                print(f"Uploading {filename} to repository root...")
+                api.upload_file(
+                    path_or_fileobj=file_path,
+                    path_in_repo=filename,
+                    repo_id=repo_id,
+                    token=token
+                )
+                uploaded_files.append(filename)
+            
+            # Upload relevant JSON config files to root (skip metadata.json and other internal files)
+            config_files_uploaded = []
+            for file_path in json_files:
+                filename = os.path.basename(file_path)
+                # Only upload important config files, skip internal metadata
+                if any(keyword in filename.lower() for keyword in ['config', 'adapter', 'lora', 'model']):
+                    print(f"Uploading {filename} to repository root...")
+                    api.upload_file(
+                        path_or_fileobj=file_path,
+                        path_in_repo=filename,
+                        repo_id=repo_id,
+                        token=token
+                    )
+                    uploaded_files.append(filename)
+                    config_files_uploaded.append(filename)
+            
+            # 2. Handle sample images
+            samples_uploaded = []
+            samples_dir = os.path.join(output_path, "samples")
+            if os.path.isdir(samples_dir):
+                print("Uploading sample images...")
+                # Create samples directory in repo
+                for filename in os.listdir(samples_dir):
+                    if filename.lower().endswith(('.jpg', '.jpeg', '.png', '.webp')):
+                        file_path = os.path.join(samples_dir, filename)
+                        repo_path = f"samples/{filename}"
+                        api.upload_file(
+                            path_or_fileobj=file_path,
+                            path_in_repo=repo_path,
+                            repo_id=repo_id,
+                            token=token
+                        )
+                        samples_uploaded.append(repo_path)
+            
+            # 3. Generate and upload README.md
+            readme_content = generate_model_card_readme(
+                repo_id=repo_id,
+                config=config,
+                model_name=model_name,
+                samples_dir=samples_dir if os.path.isdir(samples_dir) else None,
+                uploaded_files=uploaded_files
+            )
+            
+            # Create README.md file and upload to root
+            readme_path = os.path.join(temp_upload_dir, "README.md")
+            with open(readme_path, "w", encoding="utf-8") as f:
+                f.write(readme_content)
+            
+            print("Uploading README.md to repository root...")
+            api.upload_file(
+                path_or_fileobj=readme_path,
+                path_in_repo="README.md",
+                repo_id=repo_id,
+                token=token
+            )
+            
+            print(f"Model uploaded successfully to https://huggingface.co/{repo_id}")
+            print(f"Files uploaded: {len(uploaded_files)} model files, {len(samples_uploaded)} samples, README.md")
+            
+    except Exception as e:
+        print(f"Failed to upload model: {e}")
+        raise e
+
+def generate_model_card_readme(repo_id: str, config: dict, model_name: str, samples_dir: str = None, uploaded_files: list = None) -> str:
+    """Generate README.md content for the model card based on AI Toolkit's implementation"""
+    import re
+    import yaml
+    import os
+    
+    try:
+        # Extract configuration details
+        process_config = config.get("config", {}).get("process", [{}])[0]
+        model_config = process_config.get("model", {})
+        train_config = process_config.get("train", {})
+        sample_config = process_config.get("sample", {})
+        
+        # Gather model info
+        base_model = model_config.get("name_or_path", "unknown")
+        trigger_word = process_config.get("trigger_word")
+        arch = model_config.get("arch", "")
+        
+        # Determine license based on base model
+        if "FLUX.1-schnell" in base_model:
+            license_info = {"license": "apache-2.0"}
+        elif "FLUX.1-dev" in base_model:
+            license_info = {
+                "license": "other",
+                "license_name": "flux-1-dev-non-commercial-license",
+                "license_link": "https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md"
+            }
+        else:
+            license_info = {"license": "creativeml-openrail-m"}
+        
+        # Generate tags based on model architecture
+        tags = ["text-to-image"]
+        
+        if "xl" in arch.lower():
+            tags.append("stable-diffusion-xl")
+        if "flux" in arch.lower():
+            tags.append("flux")
+        if "lumina" in arch.lower():
+            tags.append("lumina2")
+        if "sd3" in arch.lower() or "v3" in arch.lower():
+            tags.append("sd3")
+        
+        # Add LoRA-specific tags
+        tags.extend(["lora", "diffusers", "template:sd-lora", "ai-toolkit"])
+        
+        # Generate widgets from sample images and prompts
+        widgets = []
+        if samples_dir and os.path.isdir(samples_dir):
+            sample_prompts = sample_config.get("samples", [])
+            if not sample_prompts:
+                # Fallback to old format
+                sample_prompts = [{"prompt": p} for p in sample_config.get("prompts", [])]
+            
+            # Get sample image files
+            sample_files = []
+            if os.path.isdir(samples_dir):
+                for filename in os.listdir(samples_dir):
+                    if filename.lower().endswith(('.jpg', '.jpeg', '.png', '.webp')):
+                        # Parse filename pattern: timestamp__steps_index.jpg
+                        match = re.search(r"__(\d+)_(\d+)\.jpg$", filename)
+                        if match:
+                            steps, index = int(match.group(1)), int(match.group(2))
+                            # Only use samples from final training step
+                            final_steps = train_config.get("steps", 1000)
+                            if steps == final_steps:
+                                sample_files.append((index, f"samples/{filename}"))
+            
+            # Sort by index and create widgets
+            sample_files.sort(key=lambda x: x[0])
+            
+            for i, prompt_obj in enumerate(sample_prompts):
+                prompt = prompt_obj.get("prompt", "") if isinstance(prompt_obj, dict) else str(prompt_obj)
+                if i < len(sample_files):
+                    _, image_path = sample_files[i]
+                    widgets.append({
+                        "text": prompt,
+                        "output": {"url": image_path}
+                    })
+        
+        # Determine torch dtype based on model
+        dtype = "torch.bfloat16" if "flux" in arch.lower() else "torch.float16"
+        
+        # Find the main safetensors file for usage example
+        main_safetensors = f"{model_name}.safetensors"
+        if uploaded_files:
+            safetensors_files = [f for f in uploaded_files if f.endswith('.safetensors')]
+            if safetensors_files:
+                main_safetensors = safetensors_files[0]
+        
+        # Construct YAML frontmatter
+        frontmatter = {
+            "tags": tags,
+            "base_model": base_model,
+            **license_info
+        }
+        
+        if widgets:
+            frontmatter["widget"] = widgets
+        
+        if trigger_word:
+            frontmatter["instance_prompt"] = trigger_word
+        
+        # Get first prompt for usage example
+        usage_prompt = trigger_word or "a beautiful landscape"
+        if widgets:
+            usage_prompt = widgets[0]["text"]
+        elif trigger_word:
+            usage_prompt = trigger_word
+        
+        # Construct README content
+        trigger_section = f"You should use \`{trigger_word}\` to trigger the image generation." if trigger_word else "No trigger words defined."
+        
+        # Build YAML frontmatter string
+        frontmatter_yaml = yaml.dump(frontmatter, default_flow_style=False, allow_unicode=True, sort_keys=False).strip()
+        
+        readme_content = f"""---
+{frontmatter_yaml}
+---
+
+# {model_name}
+
+Model trained with [AI Toolkit by Ostris](https://github.com/ostris/ai-toolkit)
+
+<Gallery />
+
+## Trigger words
+
+{trigger_section}
+
+## Download model and use it with ComfyUI, AUTOMATIC1111, SD.Next, Invoke AI, etc.
+
+Weights for this model are available in Safetensors format.
+
+[Download]({repo_id}/tree/main) them in the Files & versions tab.
+
+## Use it with the [🧨 diffusers library](https://github.com/huggingface/diffusers)
+
+\`\`\`py
+from diffusers import AutoPipelineForText2Image
+import torch
+
+pipeline = AutoPipelineForText2Image.from_pretrained('{base_model}', torch_dtype={dtype}).to('cuda')
+pipeline.load_lora_weights('{repo_id}', weight_name='{main_safetensors}')
+image = pipeline('{usage_prompt}').images[0]
+image.save("my_image.png")
+\`\`\`
+
+For more details, including weighting, merging and fusing LoRAs, check the [documentation on loading LoRAs in diffusers](https://huggingface.co/docs/diffusers/main/en/using-diffusers/loading_adapters)
+
+"""
+        return readme_content
+        
+    except Exception as e:
+        print(f"Error generating README: {e}")
+        # Fallback simple README
+        return f"""# {model_name}
+
+Model trained with [AI Toolkit by Ostris](https://github.com/ostris/ai-toolkit)
+
+## Download model
+
+Weights for this model are available in Safetensors format.
+
+[Download]({repo_id}/tree/main) them in the Files & versions tab.
+"""
+
+def main():
+    # Setup environment - token comes from HF Jobs secrets
+    if "HF_TOKEN" not in os.environ:
+        raise ValueError("HF_TOKEN environment variable not set")
+    
+    # Install system dependencies for headless operation
+    print("Installing system dependencies...")
+    try:
+        subprocess.run(["apt-get", "update"], check=True, capture_output=True)
+        subprocess.run([
+            "apt-get", "install", "-y", 
+            "libgl1-mesa-glx", 
+            "libglib2.0-0", 
+            "libsm6", 
+            "libxext6", 
+            "libxrender-dev", 
+            "libgomp1",
+            "ffmpeg"
+        ], check=True, capture_output=True)
+        print("System dependencies installed successfully")
+    except subprocess.CalledProcessError as e:
+        print(f"Failed to install system dependencies: {e}")
+        print("Continuing without system dependencies...")
+    
+    # Setup ai-toolkit
+    toolkit_dir = setup_ai_toolkit()
+    
+    # Create temporary directories
+    with tempfile.TemporaryDirectory() as temp_dir:
+        dataset_path = os.path.join(temp_dir, "dataset")
+        output_path = os.path.join(temp_dir, "output")
+        
+        # Download dataset
+        download_dataset("${datasetRepo}", dataset_path)
+        
+        # Create config
+        config = create_config(dataset_path, output_path)
+        config_path = os.path.join(temp_dir, "config.yaml")
+        
+        with open(config_path, "w") as f:
+            yaml.dump(config, f, default_flow_style=False)
+        
+        # Run training
+        print("Starting training...")
+        os.chdir(toolkit_dir)
+        
+        subprocess.run([
+            sys.executable, "run.py",
+            config_path
+        ], check=True)
+        
+        print("Training completed!")
+        
+        # Upload results
+        model_name = f"${jobConfig.config.name}-lora"
+        upload_results(output_path, model_name, "${namespace}", os.environ["HF_TOKEN"], config)
+
+if __name__ == "__main__":
+    main()
+`;
+}
+
+async function submitHFJobUV(token: string, hardware: string, scriptPath: string): Promise<string> {
+  return new Promise((resolve, reject) => {
+    // Ensure token is available
+    if (!token) {
+      reject(new Error('HF_TOKEN is required'));
+      return;
+    }
+
+    console.log('Setting up environment with HF_TOKEN for job submission');
+    console.log(`Command: hf jobs uv run --flavor ${hardware} --timeout 5h --secrets HF_TOKEN --detach ${scriptPath}`);
+    
+    // Use hf jobs uv run command with timeout and detach to get job ID
+    const childProcess = spawn('hf', [
+      'jobs', 'uv', 'run',
+      '--flavor', hardware,
+      '--timeout', '5h',
+      '--secrets', 'HF_TOKEN',
+      '--detach',
+      scriptPath
+    ], {
+      env: { 
+        ...process.env, 
+        HF_TOKEN: token 
+      }
+    });
+
+    let output = '';
+    let error = '';
+
+    childProcess.stdout.on('data', (data) => {
+      const text = data.toString();
+      output += text;
+      console.log('HF Jobs stdout:', text);
+    });
+
+    childProcess.stderr.on('data', (data) => {
+      const text = data.toString();
+      error += text;
+      console.log('HF Jobs stderr:', text);
+    });
+
+    childProcess.on('close', (code) => {
+      console.log('HF Jobs process closed with code:', code);
+      console.log('Full output:', output);
+      console.log('Full error:', error);
+      
+      if (code === 0) {
+        // With --detach flag, the output should be just the job ID
+        const fullText = (output + ' ' + error).trim();
+        
+        // Updated patterns to handle variable-length hex job IDs (16-24+ characters)
+        const jobIdPatterns = [
+          /Job started with ID:\s*([a-f0-9]{16,})/i,                    // "Job started with ID: 68b26b73767540db9fc726ac"
+          /job\s+([a-f0-9]{16,})/i,                                     // "job 68b26b73767540db9fc726ac"
+          /Job ID:\s*([a-f0-9]{16,})/i,                                 // "Job ID: 68b26b73767540db9fc726ac"
+          /created\s+job\s+([a-f0-9]{16,})/i,                          // "created job 68b26b73767540db9fc726ac"
+          /submitted.*?job\s+([a-f0-9]{16,})/i,                        // "submitted ... job 68b26b73767540db9fc726ac"
+          /https:\/\/huggingface\.co\/jobs\/[^\/]+\/([a-f0-9]{16,})/i,  // URL pattern
+          /([a-f0-9]{20,})/i,                                          // Fallback: any 20+ char hex string
+        ];
+        
+        let jobId = 'unknown';
+        
+        for (const pattern of jobIdPatterns) {
+          const match = fullText.match(pattern);
+          if (match && match[1] && match[1] !== 'started') {
+            jobId = match[1];
+            console.log(`Extracted job ID using pattern: ${pattern.toString()} -> ${jobId}`);
+            break;
+          }
+        }
+        
+        resolve(jobId);
+      } else {
+        reject(new Error(error || output || 'Failed to submit job'));
+      }
+    });
+
+    childProcess.on('error', (err) => {
+      console.error('HF Jobs process error:', err);
+      reject(new Error(`Process error: ${err.message}`));
+    });
+  });
+}
+
+async function checkHFJobStatus(token: string, jobId: string): Promise<any> {
+  return new Promise((resolve, reject) => {
+    console.log(`Checking HF Job status for: ${jobId}`);
+    
+    const childProcess = spawn('hf', [
+      'jobs', 'inspect', jobId
+    ], {
+      env: { 
+        ...process.env, 
+        HF_TOKEN: token 
+      }
+    });
+
+    let output = '';
+    let error = '';
+
+    childProcess.stdout.on('data', (data) => {
+      const text = data.toString();
+      output += text;
+    });
+
+    childProcess.stderr.on('data', (data) => {
+      const text = data.toString();
+      error += text;
+    });
+
+    childProcess.on('close', (code) => {
+      if (code === 0) {
+        try {
+          // Parse the JSON output from hf jobs inspect
+          const jobInfo = JSON.parse(output);
+          if (Array.isArray(jobInfo) && jobInfo.length > 0) {
+            const job = jobInfo[0];
+            resolve({
+              id: job.id,
+              status: job.status?.stage || 'UNKNOWN',
+              message: job.status?.message,
+              created_at: job.created_at,
+              flavor: job.flavor,
+              url: job.url,
+            });
+          } else {
+            reject(new Error('Invalid job info response'));
+          }
+        } catch (parseError: any) {
+          console.error('Failed to parse job status:', parseError, output);
+          reject(new Error('Failed to parse job status'));
+        }
+      } else {
+        reject(new Error(error || output || 'Failed to check job status'));
+      }
+    });
+
+    childProcess.on('error', (err) => {
+      console.error('HF Jobs inspect process error:', err);
+      reject(new Error(`Process error: ${err.message}`));
+    });
+  });
+}

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

@@ -0,0 +1,78 @@
+/* eslint-disable */
+import { NextRequest, NextResponse } from 'next/server';
+import fs from 'fs';
+import path from 'path';
+import { getDatasetsRoot, getTrainingFolder, getDataRoot } from '@/server/settings';
+
+export async function GET(request: NextRequest, { params }: { params: { imagePath: string } }) {
+  const { imagePath } = await params;
+  try {
+    // Decode the path
+    const filepath = decodeURIComponent(imagePath);
+
+    // Get allowed directories
+    const datasetRoot = await getDatasetsRoot();
+    const trainingRoot = await getTrainingFolder();
+    const dataRoot = await getDataRoot();
+
+    const allowedDirs = [datasetRoot, trainingRoot, dataRoot];
+
+    // Security check: Ensure path is in allowed directory
+    const isAllowed = allowedDirs.some(allowedDir => filepath.startsWith(allowedDir)) && !filepath.includes('..');
+
+    if (!isAllowed) {
+      console.warn(`Access denied: ${filepath} not in ${allowedDirs.join(', ')}`);
+      return new NextResponse('Access denied', { status: 403 });
+    }
+
+    // Check if file exists
+    if (!fs.existsSync(filepath)) {
+      console.warn(`File not found: ${filepath}`);
+      return new NextResponse('File not found', { status: 404 });
+    }
+
+    // Get file info
+    const stat = fs.statSync(filepath);
+    if (!stat.isFile()) {
+      return new NextResponse('Not a file', { status: 400 });
+    }
+
+    // Determine content type
+    const ext = path.extname(filepath).toLowerCase();
+    const contentTypeMap: { [key: string]: string } = {
+      // Images
+      '.jpg': 'image/jpeg',
+      '.jpeg': 'image/jpeg',
+      '.png': 'image/png',
+      '.gif': 'image/gif',
+      '.webp': 'image/webp',
+      '.svg': 'image/svg+xml',
+      '.bmp': 'image/bmp',
+      // Videos
+      '.mp4': 'video/mp4',
+      '.avi': 'video/x-msvideo',
+      '.mov': 'video/quicktime',
+      '.mkv': 'video/x-matroska',
+      '.wmv': 'video/x-ms-wmv',
+      '.m4v': 'video/x-m4v',
+      '.flv': 'video/x-flv'
+    };
+
+    const contentType = contentTypeMap[ext] || 'application/octet-stream';
+
+    // Read file as buffer
+    const fileBuffer = fs.readFileSync(filepath);
+
+    // Return file with appropriate headers
+    return new NextResponse(fileBuffer, {
+      headers: {
+        'Content-Type': contentType,
+        'Content-Length': String(stat.size),
+        'Cache-Control': 'public, max-age=86400',
+      },
+    });
+  } catch (error) {
+    console.error('Error serving image:', error);
+    return new NextResponse('Internal Server Error', { status: 500 });
+  }
+}

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

@@ -0,0 +1,29 @@
+import { NextResponse } from 'next/server';
+import fs from 'fs';
+import { getDatasetsRoot } from '@/server/settings';
+
+export async function POST(request: Request) {
+  try {
+    const body = await request.json();
+    const { imgPath, caption } = body;
+    let datasetsPath = await getDatasetsRoot();
+    // make sure the dataset path is in the image path
+    if (!imgPath.startsWith(datasetsPath)) {
+      return NextResponse.json({ error: 'Invalid image path' }, { status: 400 });
+    }
+
+    // if img doesnt exist, ignore
+    if (!fs.existsSync(imgPath)) {
+      return NextResponse.json({ error: 'Image does not exist' }, { status: 404 });
+    }
+
+    // check for caption
+    const captionPath = imgPath.replace(/\.[^/.]+$/, '') + '.txt';
+    // save caption to file
+    fs.writeFileSync(captionPath, caption);
+
+    return NextResponse.json({ success: true });
+  } catch (error) {
+    return NextResponse.json({ error: 'Failed to create dataset' }, { status: 500 });
+  }
+}

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

@@ -0,0 +1,34 @@
+import { NextResponse } from 'next/server';
+import fs from 'fs';
+import { getDatasetsRoot } from '@/server/settings';
+
+export async function POST(request: Request) {
+  try {
+    const body = await request.json();
+    const { imgPath } = body;
+    let datasetsPath = await getDatasetsRoot();
+    // make sure the dataset path is in the image path
+    if (!imgPath.startsWith(datasetsPath)) {
+      return NextResponse.json({ error: 'Invalid image path' }, { status: 400 });
+    }
+
+    // if img doesnt exist, ignore
+    if (!fs.existsSync(imgPath)) {
+      return NextResponse.json({ success: true });
+    }
+
+    // delete it and return success
+    fs.unlinkSync(imgPath);
+
+    // check for caption
+    const captionPath = imgPath.replace(/\.[^/.]+$/, '') + '.txt';
+    if (fs.existsSync(captionPath)) {
+      // delete caption file
+      fs.unlinkSync(captionPath);
+    }
+
+    return NextResponse.json({ success: true });
+  } catch (error) {
+    return NextResponse.json({ error: 'Failed to create dataset' }, { status: 500 });
+  }
+}

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

@@ -0,0 +1,58 @@
+// src/app/api/datasets/upload/route.ts
+import { NextRequest, NextResponse } from 'next/server';
+import { writeFile, mkdir } from 'fs/promises';
+import { join } from 'path';
+import { getDataRoot } from '@/server/settings';
+import {v4 as uuidv4} from 'uuid';
+
+export async function POST(request: NextRequest) {
+  try {
+    const dataRoot = await getDataRoot();
+    if (!dataRoot) {
+      return NextResponse.json({ error: 'Data root path not found' }, { status: 500 });
+    }
+    const imgRoot = join(dataRoot, 'images');
+
+
+    const formData = await request.formData();
+    const files = formData.getAll('files');
+
+    if (!files || files.length === 0) {
+      return NextResponse.json({ error: 'No files provided' }, { status: 400 });
+    }
+
+    // make it recursive if it doesn't exist
+    await mkdir(imgRoot, { recursive: true });
+    const savedFiles = await Promise.all(
+      files.map(async (file: any) => {
+        const bytes = await file.arrayBuffer();
+        const buffer = Buffer.from(bytes);
+
+        const extension = file.name.split('.').pop() || 'jpg';
+
+        // Clean filename and ensure it's unique
+        const fileName = `${uuidv4()}`; // Use UUID for unique file names
+        const filePath = join(imgRoot, `${fileName}.${extension}`);
+
+        await writeFile(filePath, buffer);
+        return filePath;
+      }),
+    );
+
+    return NextResponse.json({
+      message: 'Files uploaded successfully',
+      files: savedFiles,
+    });
+  } catch (error) {
+    console.error('Upload error:', error);
+    return NextResponse.json({ error: 'Error uploading files' }, { status: 500 });
+  }
+}
+
+// Increase payload size limit (default is 4mb)
+export const config = {
+  api: {
+    bodyParser: false,
+    responseLimit: '50mb',
+  },
+};

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

@@ -0,0 +1,32 @@
+import { NextRequest, NextResponse } from 'next/server';
+import { PrismaClient } from '@prisma/client';
+import { getTrainingFolder } from '@/server/settings';
+import path from 'path';
+import fs from 'fs';
+
+const prisma = new PrismaClient();
+
+export async function GET(request: NextRequest, { params }: { params: { jobID: string } }) {
+  const { jobID } = await params;
+
+  const job = await prisma.job.findUnique({
+    where: { id: jobID },
+  });
+
+  if (!job) {
+    return NextResponse.json({ error: 'Job not found' }, { status: 404 });
+  }
+
+  const trainingRoot = await getTrainingFolder();
+  const trainingFolder = path.join(trainingRoot, job.name);
+
+  if (fs.existsSync(trainingFolder)) {
+    fs.rmdirSync(trainingFolder, { recursive: true });
+  }
+
+  await prisma.job.delete({
+    where: { id: jobID },
+  });
+
+  return NextResponse.json(job);
+}

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

@@ -0,0 +1,48 @@
+import { NextRequest, NextResponse } from 'next/server';
+import { PrismaClient } from '@prisma/client';
+import path from 'path';
+import fs from 'fs';
+import { getTrainingFolder } from '@/server/settings';
+
+const prisma = new PrismaClient();
+
+export async function GET(request: NextRequest, { params }: { params: { jobID: string } }) {
+  const { jobID } = await params;
+
+  const job = await prisma.job.findUnique({
+    where: { id: jobID },
+  });
+
+  if (!job) {
+    return NextResponse.json({ error: 'Job not found' }, { status: 404 });
+  }
+
+  const trainingFolder = await getTrainingFolder();
+  const jobFolder = path.join(trainingFolder, job.name);
+
+  if (!fs.existsSync(jobFolder)) {
+    return NextResponse.json({ files: [] });
+  }
+
+  // find all safetensors files in the job folder
+  let files = fs
+    .readdirSync(jobFolder)
+    .filter(file => {
+      return file.endsWith('.safetensors');
+    })
+    .map(file => {
+      return path.join(jobFolder, file);
+    })
+    .sort();
+
+  // get the file size for each file
+  const fileObjects = files.map(file => {
+    const stats = fs.statSync(file);
+    return {
+      path: file,
+      size: stats.size,
+    };
+  });
+
+  return NextResponse.json({ files: fileObjects });
+}

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

@@ -0,0 +1,35 @@
+import { NextRequest, NextResponse } from 'next/server';
+import { PrismaClient } from '@prisma/client';
+import path from 'path';
+import fs from 'fs';
+import { getTrainingFolder } from '@/server/settings';
+
+const prisma = new PrismaClient();
+
+export async function GET(request: NextRequest, { params }: { params: { jobID: string } }) {
+  const { jobID } = await params;
+
+  const job = await prisma.job.findUnique({
+    where: { id: jobID },
+  });
+
+  if (!job) {
+    return NextResponse.json({ error: 'Job not found' }, { status: 404 });
+  }
+
+  const trainingFolder = await getTrainingFolder();
+  const jobFolder = path.join(trainingFolder, job.name);
+  const logPath = path.join(jobFolder, 'log.txt');
+
+  if (!fs.existsSync(logPath)) {
+    return NextResponse.json({ log: '' });
+  }
+  let log = '';
+  try {
+    log = fs.readFileSync(logPath, 'utf-8');
+  } catch (error) {
+    console.error('Error reading log file:', error);
+    log = 'Error reading log file';
+  }
+  return NextResponse.json({ log: log });
+}

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

@@ -0,0 +1,40 @@
+import { NextRequest, NextResponse } from 'next/server';
+import { PrismaClient } from '@prisma/client';
+import path from 'path';
+import fs from 'fs';
+import { getTrainingFolder } from '@/server/settings';
+
+const prisma = new PrismaClient();
+
+export async function GET(request: NextRequest, { params }: { params: { jobID: string } }) {
+  const { jobID } = await params;
+
+  const job = await prisma.job.findUnique({
+    where: { id: jobID },
+  });
+
+  if (!job) {
+    return NextResponse.json({ error: 'Job not found' }, { status: 404 });
+  }
+
+  // setup the training
+  const trainingFolder = await getTrainingFolder();
+
+  const samplesFolder = path.join(trainingFolder, job.name, 'samples');
+  if (!fs.existsSync(samplesFolder)) {
+    return NextResponse.json({ samples: [] });
+  }
+
+  // find all img (png, jpg, jpeg) files in the samples folder
+  const samples = fs
+    .readdirSync(samplesFolder)
+    .filter(file => {
+      return file.endsWith('.png') || file.endsWith('.jpg') || file.endsWith('.jpeg') || file.endsWith('.webp');
+    })
+    .map(file => {
+      return path.join(samplesFolder, file);
+    })
+    .sort();
+
+  return NextResponse.json({ samples });
+}

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

@@ -0,0 +1,215 @@
+import { NextRequest, NextResponse } from 'next/server';
+import { PrismaClient } from '@prisma/client';
+import { TOOLKIT_ROOT } from '@/paths';
+import { spawn } from 'child_process';
+import path from 'path';
+import fs from 'fs';
+import os from 'os';
+import { getTrainingFolder, getHFToken } from '@/server/settings';
+const isWindows = process.platform === 'win32';
+
+const prisma = new PrismaClient();
+
+export async function GET(request: NextRequest, { params }: { params: { jobID: string } }) {
+  const { jobID } = await params;
+
+  const job = await prisma.job.findUnique({
+    where: { id: jobID },
+  });
+
+  if (!job) {
+    return NextResponse.json({ error: 'Job not found' }, { status: 404 });
+  }
+
+  // update job status to 'running'
+  await prisma.job.update({
+    where: { id: jobID },
+    data: {
+      status: 'running',
+      stop: false,
+      info: 'Starting job...',
+    },
+  });
+
+  // setup the training
+  const trainingRoot = await getTrainingFolder();
+
+  const trainingFolder = path.join(trainingRoot, job.name);
+  if (!fs.existsSync(trainingFolder)) {
+    fs.mkdirSync(trainingFolder, { recursive: true });
+  }
+
+  // make the config file
+  const configPath = path.join(trainingFolder, '.job_config.json');
+
+  //log to path
+  const logPath = path.join(trainingFolder, 'log.txt');
+
+  try {
+    // if the log path exists, move it to a folder called logs and rename it {num}_log.txt, looking for the highest num
+    // if the log path does not exist, create it
+    if (fs.existsSync(logPath)) {
+      const logsFolder = path.join(trainingFolder, 'logs');
+      if (!fs.existsSync(logsFolder)) {
+        fs.mkdirSync(logsFolder, { recursive: true });
+      }
+
+      let num = 0;
+      while (fs.existsSync(path.join(logsFolder, `${num}_log.txt`))) {
+        num++;
+      }
+
+      fs.renameSync(logPath, path.join(logsFolder, `${num}_log.txt`));
+    }
+  } catch (e) {
+    console.error('Error moving log file:', e);
+  }
+
+  // update the config dataset path
+  const jobConfig = JSON.parse(job.job_config);
+  jobConfig.config.process[0].sqlite_db_path = path.join(TOOLKIT_ROOT, 'aitk_db.db');
+
+  // write the config file
+  fs.writeFileSync(configPath, JSON.stringify(jobConfig, null, 2));
+
+  let pythonPath = 'python';
+  // use .venv or venv if it exists
+  if (fs.existsSync(path.join(TOOLKIT_ROOT, '.venv'))) {
+    if (isWindows) {
+      pythonPath = path.join(TOOLKIT_ROOT, '.venv', 'Scripts', 'python.exe');
+    } else {
+      pythonPath = path.join(TOOLKIT_ROOT, '.venv', 'bin', 'python');
+    }
+  } else if (fs.existsSync(path.join(TOOLKIT_ROOT, 'venv'))) {
+    if (isWindows) {
+      pythonPath = path.join(TOOLKIT_ROOT, 'venv', 'Scripts', 'python.exe');
+    } else {
+      pythonPath = path.join(TOOLKIT_ROOT, 'venv', 'bin', 'python');
+    }
+  }
+
+  const runFilePath = path.join(TOOLKIT_ROOT, 'run.py');
+  if (!fs.existsSync(runFilePath)) {
+    return NextResponse.json({ error: 'run.py not found' }, { status: 500 });
+  }
+
+  const additionalEnv: any = {
+    AITK_JOB_ID: jobID,
+    CUDA_VISIBLE_DEVICES: `${job.gpu_ids}`,
+    IS_AI_TOOLKIT_UI: '1'
+  };
+
+  // HF_TOKEN
+  const hfToken = await getHFToken();
+  if (hfToken && hfToken.trim() !== '') {
+    additionalEnv.HF_TOKEN = hfToken;
+  }
+
+  // Add the --log argument to the command
+  const args = [runFilePath, configPath, '--log', logPath];
+
+  try {
+    let subprocess;
+
+    if (isWindows) {
+      // For Windows, use 'cmd.exe' to open a new command window
+      subprocess = spawn('cmd.exe', ['/c', 'start', 'cmd.exe', '/k', pythonPath, ...args], {
+        env: {
+          ...process.env,
+          ...additionalEnv,
+        },
+        cwd: TOOLKIT_ROOT,
+        windowsHide: false,
+      });
+    } else {
+      // For non-Windows platforms
+      subprocess = spawn(pythonPath, args, {
+        detached: true,
+        stdio: ['ignore', 'pipe', 'pipe'], // Changed from 'ignore' to capture output
+        env: {
+          ...process.env,
+          ...additionalEnv,
+        },
+        cwd: TOOLKIT_ROOT,
+      });
+    }
+
+    // Start monitoring in the background without blocking the response
+    const monitorProcess = async () => {
+      const startTime = Date.now();
+      let errorOutput = '';
+      let stdoutput = '';
+
+      if (subprocess.stderr) {
+        subprocess.stderr.on('data', data => {
+          errorOutput += data.toString();
+        });
+        subprocess.stdout.on('data', data => {
+          stdoutput += data.toString();
+          // truncate to only get the last 500 characters
+          if (stdoutput.length > 500) {
+            stdoutput = stdoutput.substring(stdoutput.length - 500);
+          }
+        });
+      }
+
+      subprocess.on('exit', async code => {
+        const currentTime = Date.now();
+        const duration = (currentTime - startTime) / 1000;
+        console.log(`Job ${jobID} exited with code ${code} after ${duration} seconds.`);
+        // wait for 5 seconds to give it time to stop itself. It id still has a status of running in the db, update it to stopped
+        await new Promise(resolve => setTimeout(resolve, 5000));
+        const updatedJob = await prisma.job.findUnique({
+          where: { id: jobID },
+        });
+        if (updatedJob?.status === 'running') {
+          let errorString = errorOutput;
+          if (errorString.trim() === '') {
+            errorString = stdoutput;
+          }
+          await prisma.job.update({
+            where: { id: jobID },
+            data: {
+              status: 'error',
+              info: `Error launching job: ${errorString.substring(0, 500)}`,
+            },
+          });
+        }
+      });
+
+      // Wait 30 seconds before releasing the process
+      await new Promise(resolve => setTimeout(resolve, 30000));
+      // Detach the process for non-Windows systems
+      if (!isWindows && subprocess.unref) {
+        subprocess.unref();
+      }
+    };
+
+    // Start the monitoring without awaiting it
+    monitorProcess().catch(err => {
+      console.error(`Error in process monitoring for job ${jobID}:`, err);
+    });
+
+    // Return the response immediately
+    return NextResponse.json(job);
+  } catch (error: any) {
+    // Handle any exceptions during process launch
+    console.error('Error launching process:', error);
+
+    await prisma.job.update({
+      where: { id: jobID },
+      data: {
+        status: 'error',
+        info: `Error launching job: ${error?.message || 'Unknown error'}`,
+      },
+    });
+
+    return NextResponse.json(
+      {
+        error: 'Failed to launch job process',
+        details: error?.message || 'Unknown error',
+      },
+      { status: 500 },
+    );
+  }
+}

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

@@ -0,0 +1,23 @@
+import { NextRequest, NextResponse } from 'next/server';
+import { PrismaClient } from '@prisma/client';
+
+const prisma = new PrismaClient();
+
+export async function GET(request: NextRequest, { params }: { params: { jobID: string } }) {
+  const { jobID } = await params;
+
+  const job = await prisma.job.findUnique({
+    where: { id: jobID },
+  });
+
+  // update job status to 'running'
+  await prisma.job.update({
+    where: { id: jobID },
+    data: {
+      stop: true,
+      info: 'Stopping job...',
+    },
+  });
+
+  return NextResponse.json(job);
+}

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

@@ -0,0 +1,67 @@
+import { NextResponse } from 'next/server';
+import { PrismaClient } from '@prisma/client';
+
+const prisma = new PrismaClient();
+
+export async function GET(request: Request) {
+  const { searchParams } = new URL(request.url);
+  const id = searchParams.get('id');
+
+  try {
+    if (id) {
+      const job = await prisma.job.findUnique({
+        where: { id },
+      });
+      return NextResponse.json(job);
+    }
+
+    const jobs = await prisma.job.findMany({
+      orderBy: { created_at: 'desc' },
+    });
+    return NextResponse.json({ jobs: jobs });
+  } catch (error) {
+    console.error(error);
+    return NextResponse.json({ error: 'Failed to fetch training data' }, { status: 500 });
+  }
+}
+
+export async function POST(request: Request) {
+  try {
+    const body = await request.json();
+    const { id, name, job_config, gpu_ids } = body;
+
+    // Ensure gpu_ids is never null/undefined - provide default value
+    const safeGpuIds = gpu_ids || '0';
+
+    if (id) {
+      // Update existing training
+      const training = await prisma.job.update({
+        where: { id },
+        data: {
+          name,
+          gpu_ids: safeGpuIds,
+          job_config: JSON.stringify(job_config),
+        },
+      });
+      return NextResponse.json(training);
+    } else {
+      // Create new training
+      const training = await prisma.job.create({
+        data: {
+          name,
+          gpu_ids: safeGpuIds,
+          job_config: JSON.stringify(job_config),
+        },
+      });
+      return NextResponse.json(training);
+    }
+  } catch (error: any) {
+    if (error.code === 'P2002') {
+      // Handle unique constraint violation, 409=Conflict
+      return NextResponse.json({ error: 'Job name already exists' }, { status: 409 });
+    }
+    console.error(error);
+    // Handle other errors
+    return NextResponse.json({ error: 'Failed to save training data' }, { status: 500 });
+  }
+}

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

@@ -0,0 +1,59 @@
+import { NextResponse } from 'next/server';
+import { PrismaClient } from '@prisma/client';
+import { defaultTrainFolder, defaultDatasetsFolder } from '@/paths';
+import { flushCache } from '@/server/settings';
+
+const prisma = new PrismaClient();
+
+export async function GET() {
+  try {
+    const settings = await prisma.settings.findMany();
+    const settingsObject = settings.reduce((acc: any, setting) => {
+      acc[setting.key] = setting.value;
+      return acc;
+    }, {});
+    // if TRAINING_FOLDER is not set, use default
+    if (!settingsObject.TRAINING_FOLDER || settingsObject.TRAINING_FOLDER === '') {
+      settingsObject.TRAINING_FOLDER = defaultTrainFolder;
+    }
+    // if DATASETS_FOLDER is not set, use default
+    if (!settingsObject.DATASETS_FOLDER || settingsObject.DATASETS_FOLDER === '') {
+      settingsObject.DATASETS_FOLDER = defaultDatasetsFolder;
+    }
+    return NextResponse.json(settingsObject);
+  } catch (error) {
+    return NextResponse.json({ error: 'Failed to fetch settings' }, { status: 500 });
+  }
+}
+
+export async function POST(request: Request) {
+  try {
+    const body = await request.json();
+    const { HF_TOKEN, TRAINING_FOLDER, DATASETS_FOLDER } = body;
+
+    // Upsert both settings
+    await Promise.all([
+      prisma.settings.upsert({
+        where: { key: 'HF_TOKEN' },
+        update: { value: HF_TOKEN },
+        create: { key: 'HF_TOKEN', value: HF_TOKEN },
+      }),
+      prisma.settings.upsert({
+        where: { key: 'TRAINING_FOLDER' },
+        update: { value: TRAINING_FOLDER },
+        create: { key: 'TRAINING_FOLDER', value: TRAINING_FOLDER },
+      }),
+      prisma.settings.upsert({
+        where: { key: 'DATASETS_FOLDER' },
+        update: { value: DATASETS_FOLDER },
+        create: { key: 'DATASETS_FOLDER', value: DATASETS_FOLDER },
+      }),
+    ]);
+
+    flushCache();
+
+    return NextResponse.json({ success: true });
+  } catch (error) {
+    return NextResponse.json({ error: 'Failed to update settings' }, { status: 500 });
+  }
+}

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

@@ -0,0 +1,78 @@
+/* eslint-disable */
+import { NextRequest, NextResponse } from 'next/server';
+import fs from 'fs';
+import fsp from 'fs/promises';
+import path from 'path';
+import archiver from 'archiver';
+import { getTrainingFolder } from '@/server/settings';
+
+export const runtime = 'nodejs'; // ensure Node APIs are available
+export const dynamic = 'force-dynamic'; // long-running, non-cached
+
+type PostBody = {
+  zipTarget: 'samples'; //only samples for now
+  jobName: string;
+};
+
+async function resolveSafe(p: string) {
+  // resolve symlinks + normalize
+  return await fsp.realpath(p);
+}
+
+export async function POST(request: NextRequest) {
+  try {
+    const body = (await request.json()) as PostBody;
+    if (!body || !body.jobName) {
+      return NextResponse.json({ error: 'jobName is required' }, { status: 400 });
+    }
+
+    const trainingRoot = await resolveSafe(await getTrainingFolder());
+    const folderPath = await resolveSafe(path.join(trainingRoot, body.jobName, 'samples'));
+    const outputPath = path.resolve(trainingRoot, body.jobName, 'samples.zip');
+
+    // Must be a directory
+    let stat: fs.Stats;
+    try {
+      stat = await fsp.stat(folderPath);
+    } catch {
+      return new NextResponse('Folder not found', { status: 404 });
+    }
+    if (!stat.isDirectory()) {
+      return new NextResponse('Not a directory', { status: 400 });
+    }
+
+    // delete current one if it exists
+    if (fs.existsSync(outputPath)) {
+      await fsp.unlink(outputPath);
+    }
+
+    // Create write stream & archive
+    await new Promise<void>((resolve, reject) => {
+      const output = fs.createWriteStream(outputPath);
+      const archive = archiver('zip', { zlib: { level: 9 } });
+
+      output.on('close', () => resolve());
+      output.on('error', reject);
+      archive.on('error', reject);
+
+      archive.pipe(output);
+
+      // Add the directory contents (place them under the folder's base name in the zip)
+      const rootName = path.basename(folderPath);
+      archive.directory(folderPath, rootName);
+
+      archive.finalize().catch(reject);
+    });
+
+    // Return the absolute path so your existing /api/files/[...filePath] can serve it
+    // Example download URL (client-side): `/api/files/${encodeURIComponent(resolvedOutPath)}`
+    return NextResponse.json({
+      ok: true,
+      zipPath: outputPath,
+      fileName: path.basename(outputPath),
+    });
+  } catch (err) {
+    console.error('Zip error:', err);
+    return new NextResponse('Internal Server Error', { status: 500 });
+  }
+}

ui/src/app/dashboard/page.tsx

@@ -0,0 +1,54 @@
+'use client';
+
+import GpuMonitor from '@/components/GPUMonitor';
+import JobsTable from '@/components/JobsTable';
+import { TopBar, MainContent } from '@/components/layout';
+import Link from 'next/link';
+import { useAuth } from '@/contexts/AuthContext';
+import HFLoginButton from '@/components/HFLoginButton';
+
+export default function Dashboard() {
+  const { status: authStatus, namespace } = useAuth();
+  const isAuthenticated = authStatus === 'authenticated';
+
+  return (
+    <>
+      <TopBar>
+        <div>
+          <h1 className="text-lg">Dashboard</h1>
+        </div>
+        <div className="flex-1 flex items-center justify-end gap-3 pr-2 text-sm text-gray-400">
+          {isAuthenticated ? (
+            <span>Welcome, {namespace || 'user'}</span>
+          ) : (
+            <>
+              <span>Welcome, Guest</span>
+              <HFLoginButton size="sm" />
+              <Link href="/settings" className="text-xs text-blue-400 hover:text-blue-300">
+                Settings
+              </Link>
+            </>
+          )}
+        </div>
+      </TopBar>
+      <MainContent>
+        <GpuMonitor />
+        <div className="w-full mt-4">
+          <div className="flex justify-between items-center mb-2">
+            <h1 className="text-md">Active Jobs</h1>
+            <div className="text-xs text-gray-500">
+              <Link href="/jobs">View All</Link>
+            </div>
+          </div>
+          {isAuthenticated ? (
+            <JobsTable onlyActive />
+          ) : (
+            <div className="border border-gray-800 rounded-lg p-6 bg-gray-900 text-gray-400 text-sm">
+              Sign in with Hugging Face or add an access token in Settings to view and manage jobs.
+            </div>
+          )}
+        </div>
+      </MainContent>
+    </>
+  );
+}

ui/src/app/datasets/[datasetName]/page.tsx

@@ -0,0 +1,190 @@
+'use client';
+
+import { useEffect, useState, use, useMemo } from 'react';
+import { LuImageOff, LuLoader, LuBan } from 'react-icons/lu';
+import { FaChevronLeft } from 'react-icons/fa';
+import DatasetImageCard from '@/components/DatasetImageCard';
+import { Button } from '@headlessui/react';
+import AddImagesModal, { openImagesModal } from '@/components/AddImagesModal';
+import { TopBar, MainContent } from '@/components/layout';
+import { apiClient } from '@/utils/api';
+import FullscreenDropOverlay from '@/components/FullscreenDropOverlay';
+import { useRouter } from 'next/navigation';
+import { usingBrowserDb } from '@/utils/env';
+import { hasUserDataset } from '@/utils/storage/datasetStorage';
+import { useAuth } from '@/contexts/AuthContext';
+import HFLoginButton from '@/components/HFLoginButton';
+import Link from 'next/link';
+
+export default function DatasetPage({ params }: { params: { datasetName: string } }) {
+  const [imgList, setImgList] = useState<{ img_path: string }[]>([]);
+  const usableParams = use(params as any) as { datasetName: string };
+  const datasetName = usableParams.datasetName;
+  const [status, setStatus] = useState<'idle' | 'loading' | 'success' | 'error'>('idle');
+  const router = useRouter();
+  const { status: authStatus } = useAuth();
+  const isAuthenticated = authStatus === 'authenticated';
+  const hasDatasetEntry = !usingBrowserDb || hasUserDataset(datasetName);
+  const allowAccess = hasDatasetEntry && isAuthenticated;
+
+  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));
+        setImgList(data.images);
+        setStatus('success');
+      })
+      .catch(error => {
+        console.error('Error fetching images:', error);
+        setStatus('error');
+      });
+  };
+  useEffect(() => {
+    if (!datasetName) {
+      return;
+    }
+
+    if (!isAuthenticated) {
+      return;
+    }
+
+    if (!hasDatasetEntry) {
+      setImgList([]);
+      setStatus('error');
+      router.replace('/datasets');
+      return;
+    }
+
+    refreshImageList(datasetName);
+  }, [datasetName, hasDatasetEntry, isAuthenticated, router]);
+
+  if (!allowAccess) {
+    return (
+      <>
+        <TopBar>
+          <div>
+            <Button className="text-gray-500 dark:text-gray-300 px-3 mt-1" onClick={() => history.back()}>
+              <FaChevronLeft />
+            </Button>
+          </div>
+          <div>
+            <h1 className="text-lg">Dataset: {datasetName}</h1>
+          </div>
+          <div className="flex-1"></div>
+        </TopBar>
+        <MainContent>
+          <div className="border border-gray-800 rounded-lg p-6 bg-gray-900 text-gray-400 text-sm flex flex-col gap-4">
+            <p>You need to sign in with Hugging Face or provide a valid token to view this dataset.</p>
+            <div className="flex items-center gap-3">
+              <HFLoginButton size="sm" />
+              <Link href="/settings" className="text-xs text-blue-400 hover:text-blue-300">
+                Manage authentication in Settings
+              </Link>
+            </div>
+          </div>
+        </MainContent>
+      </>
+    );
+  }
+
+  const PageInfoContent = useMemo(() => {
+    let icon = null;
+    let text = '';
+    let subtitle = '';
+    let showIt = false;
+    let bgColor = '';
+    let textColor = '';
+    let iconColor = '';
+
+    if (status == 'loading') {
+      icon = <LuLoader className="animate-spin w-8 h-8" />;
+      text = 'Loading Images';
+      subtitle = 'Please wait while we fetch your dataset images...';
+      showIt = true;
+      bgColor = 'bg-gray-50 dark:bg-gray-800/50';
+      textColor = 'text-gray-900 dark:text-gray-100';
+      iconColor = 'text-gray-500 dark:text-gray-400';
+    }
+    if (status == 'error') {
+      icon = <LuBan className="w-8 h-8" />;
+      text = 'Error Loading Images';
+      subtitle = 'There was a problem fetching the images. Please try refreshing the page.';
+      showIt = true;
+      bgColor = 'bg-red-50 dark:bg-red-950/20';
+      textColor = 'text-red-900 dark:text-red-100';
+      iconColor = 'text-red-600 dark:text-red-400';
+    }
+    if (status == 'success' && imgList.length === 0) {
+      icon = <LuImageOff className="w-8 h-8" />;
+      text = 'No Images Found';
+      subtitle = 'This dataset is empty. Click "Add Images" to get started.';
+      showIt = true;
+      bgColor = 'bg-gray-50 dark:bg-gray-800/50';
+      textColor = 'text-gray-900 dark:text-gray-100';
+      iconColor = 'text-gray-500 dark:text-gray-400';
+    }
+
+    if (!showIt) return null;
+
+    return (
+      <div
+        className={`mt-10 flex flex-col items-center justify-center py-16 px-8 rounded-xl border-2 border-gray-700 border-dashed ${bgColor} ${textColor} mx-auto max-w-md text-center`}
+      >
+        <div className={`${iconColor} mb-4`}>{icon}</div>
+        <h3 className="text-lg font-semibold mb-2">{text}</h3>
+        <p className="text-sm opacity-75 leading-relaxed">{subtitle}</p>
+      </div>
+    );
+  }, [status, imgList.length]);
+
+  return (
+    <>
+      {/* Fixed top bar */}
+      <TopBar>
+        <div>
+          <Button className="text-gray-500 dark:text-gray-300 px-3 mt-1" onClick={() => history.back()}>
+            <FaChevronLeft />
+          </Button>
+        </div>
+        <div>
+          <h1 className="text-lg">Dataset: {datasetName}</h1>
+        </div>
+        <div className="flex-1"></div>
+        <div>
+          <Button
+            className="text-gray-200 bg-slate-600 px-3 py-1 rounded-md"
+            onClick={() => openImagesModal(datasetName, () => refreshImageList(datasetName))}
+          >
+            Add Images
+          </Button>
+        </div>
+      </TopBar>
+      <MainContent>
+        {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"
+                imageUrl={img.img_path}
+                onDelete={() => refreshImageList(datasetName)}
+              />
+            ))}
+          </div>
+        )}
+      </MainContent>
+      <AddImagesModal />
+      <FullscreenDropOverlay
+        datasetName={datasetName}
+        onComplete={() => refreshImageList(datasetName)}
+      />
+    </>
+  );
+}

ui/src/app/datasets/page.tsx

@@ -0,0 +1,217 @@
+'use client';
+
+import { useState } from 'react';
+import { Modal } from '@/components/Modal';
+import Link from 'next/link';
+import { TextInput } from '@/components/formInputs';
+import useDatasetList from '@/hooks/useDatasetList';
+import { Button } from '@headlessui/react';
+import { FaRegTrashAlt } from 'react-icons/fa';
+import { openConfirm } from '@/components/ConfirmModal';
+import { TopBar, MainContent } from '@/components/layout';
+import UniversalTable, { TableColumn } from '@/components/UniversalTable';
+import { apiClient } from '@/utils/api';
+import { useRouter } from 'next/navigation';
+import { usingBrowserDb } from '@/utils/env';
+import { addUserDataset, removeUserDataset } from '@/utils/storage/datasetStorage';
+import { useAuth } from '@/contexts/AuthContext';
+import HFLoginButton from '@/components/HFLoginButton';
+
+export default function Datasets() {
+  const router = useRouter();
+  const { datasets, status, refreshDatasets } = useDatasetList();
+  const [newDatasetName, setNewDatasetName] = useState('');
+  const [isNewDatasetModalOpen, setIsNewDatasetModalOpen] = useState(false);
+  const { status: authStatus } = useAuth();
+  const isAuthenticated = authStatus === 'authenticated';
+
+  // Transform datasets array into rows with objects
+  const tableRows = datasets.map(dataset => ({
+    name: dataset,
+    actions: dataset, // Pass full dataset name for actions
+  }));
+
+  const columns: TableColumn[] = [
+    {
+      title: 'Dataset Name',
+      key: 'name',
+      render: row => (
+        <Link href={`/datasets/${row.name}`} className="text-gray-200 hover:text-gray-100">
+          {row.name}
+        </Link>
+      ),
+    },
+    {
+      title: 'Actions',
+      key: 'actions',
+      className: 'w-20 text-right',
+      render: row => (
+        <button
+          className="text-gray-200 hover:bg-red-600 p-2 rounded-full transition-colors"
+          onClick={() => handleDeleteDataset(row.name)}
+        >
+          <FaRegTrashAlt />
+        </button>
+      ),
+    },
+  ];
+
+  const handleDeleteDataset = (datasetName: string) => {
+    openConfirm({
+      title: 'Delete Dataset',
+      message: `Are you sure you want to delete the dataset "${datasetName}"? This action cannot be undone.`,
+      type: 'warning',
+      confirmText: 'Delete',
+      onConfirm: () => {
+        apiClient
+          .post('/api/datasets/delete', { name: datasetName })
+          .then(() => {
+            console.log('Dataset deleted:', datasetName);
+            if (usingBrowserDb) {
+              removeUserDataset(datasetName);
+            }
+            refreshDatasets();
+          })
+          .catch(error => {
+            console.error('Error deleting dataset:', error);
+          });
+      },
+    });
+  };
+
+  const handleCreateDataset = async (e: React.FormEvent) => {
+    e.preventDefault();
+    if (!isAuthenticated) {
+      return;
+    }
+    try {
+      const data = await apiClient.post('/api/datasets/create', { name: newDatasetName }).then(res => res.data);
+      console.log('New dataset created:', data);
+      if (usingBrowserDb && data?.name) {
+        addUserDataset(data.name, data?.path || '');
+      }
+      refreshDatasets();
+      setNewDatasetName('');
+      setIsNewDatasetModalOpen(false);
+    } catch (error) {
+      console.error('Error creating new dataset:', error);
+    }
+  };
+
+  const openNewDatasetModal = () => {
+    if (!isAuthenticated) {
+      return;
+    }
+    openConfirm({
+      title: 'New Dataset',
+      message: 'Enter the name of the new dataset:',
+      type: 'info',
+      confirmText: 'Create',
+      inputTitle: 'Dataset Name',
+      onConfirm: async (name?: string) => {
+        if (!name) {
+          console.error('Dataset name is required.');
+          return;
+        }
+        if (!isAuthenticated) {
+          return;
+        }
+        try {
+          const data = await apiClient.post('/api/datasets/create', { name }).then(res => res.data);
+          console.log('New dataset created:', data);
+          if (usingBrowserDb && data?.name) {
+            addUserDataset(data.name, data?.path || '');
+          }
+          if (data.name) {
+            router.push(`/datasets/${data.name}`);
+          } else {
+            refreshDatasets();
+          }
+        } catch (error) {
+          console.error('Error creating new dataset:', error);
+        }
+      },
+    });
+  };
+
+  return (
+    <>
+      <TopBar>
+        <div>
+          <h1 className="text-2xl font-semibold text-gray-100">Datasets</h1>
+        </div>
+        <div className="flex-1"></div>
+        <div>
+          {isAuthenticated ? (
+            <Button
+              className="text-gray-200 bg-slate-600 px-4 py-2 rounded-md hover:bg-slate-500 transition-colors"
+              onClick={() => openNewDatasetModal()}
+            >
+              New Dataset
+            </Button>
+          ) : (
+            <span className="text-gray-600 bg-gray-900 px-3 py-1 rounded-md border border-gray-800">
+              Sign in to add datasets
+            </span>
+          )}
+        </div>
+      </TopBar>
+
+      <MainContent>
+        {isAuthenticated ? (
+          <UniversalTable
+            columns={columns}
+            rows={tableRows}
+            isLoading={status === 'loading'}
+            onRefresh={refreshDatasets}
+          />
+        ) : (
+          <div className="border border-gray-800 rounded-lg p-6 bg-gray-900 text-gray-400 text-sm flex flex-col gap-4">
+            <p>Sign in with Hugging Face or add an access token to manage datasets.</p>
+            <div className="flex items-center gap-3">
+              <HFLoginButton size="sm" />
+              <Link href="/settings" className="text-xs text-blue-400 hover:text-blue-300">
+                Manage authentication in Settings
+              </Link>
+            </div>
+          </div>
+        )}
+      </MainContent>
+
+      <Modal
+        isOpen={isNewDatasetModalOpen}
+        onClose={() => setIsNewDatasetModalOpen(false)}
+        title="New Dataset"
+        size="md"
+      >
+        <div className="space-y-4 text-gray-200">
+          <form onSubmit={handleCreateDataset}>
+            <div className="text-sm text-gray-400">
+              This will create a new folder with the name below in your dataset folder.
+            </div>
+            <div className="mt-4">
+              <TextInput label="Dataset Name" value={newDatasetName} onChange={value => setNewDatasetName(value)} />
+            </div>
+
+            <div className="mt-6 flex justify-end space-x-3">
+              <button
+                type="button"
+                className="rounded-md bg-gray-700 px-4 py-2 text-gray-200 hover:bg-gray-600 focus:outline-none focus:ring-2 focus:ring-gray-500"
+                onClick={() => setIsNewDatasetModalOpen(false)}
+              >
+                Cancel
+              </button>
+              <button
+                type="submit"
+                className="rounded-md bg-blue-600 px-4 py-2 text-white hover:bg-blue-700 focus:outline-none focus:ring-2 focus:ring-blue-500 disabled:opacity-50 disabled:cursor-not-allowed"
+                disabled={!isAuthenticated}
+              >
+                Confirm
+              </button>
+            </div>
+          </form>
+        </div>
+      </Modal>
+    </>
+  );
+}