Spaces:

multimodalart
/

ai-toolkit

Running on CPU Upgrade

App Files Files Community

multimodalart HF Staff commited on Sep 17

Commit

1125d73

verified ·

1 Parent(s): 2ed6d27

Upload 99 files

Browse files

This view is limited to 50 files because it contains too many changes. See raw diff

Files changed (50) hide show

src/app/api/auth/hf/callback/route.ts +112 -0
src/app/api/auth/hf/login/route.ts +36 -0
src/app/api/auth/hf/validate/route.ts +22 -0
src/app/api/auth/route.ts +6 -0
src/app/api/caption/get/route.ts +46 -0
src/app/api/datasets/create/route.tsx +25 -0
src/app/api/datasets/delete/route.tsx +24 -0
src/app/api/datasets/list/route.ts +25 -0
src/app/api/datasets/listImages/route.ts +61 -0
src/app/api/datasets/upload/route.ts +57 -0
src/app/api/files/[...filePath]/route.ts +116 -0
src/app/api/gpu/route.ts +121 -0
src/app/api/hf-hub/route.ts +165 -0
src/app/api/hf-jobs/route.ts +761 -0
src/app/api/img/[...imagePath]/route.ts +78 -0
src/app/api/img/caption/route.ts +29 -0
src/app/api/img/delete/route.ts +34 -0
src/app/api/img/upload/route.ts +58 -0
src/app/api/jobs/[jobID]/delete/route.ts +32 -0
src/app/api/jobs/[jobID]/files/route.ts +48 -0
src/app/api/jobs/[jobID]/log/route.ts +35 -0
src/app/api/jobs/[jobID]/samples/route.ts +40 -0
src/app/api/jobs/[jobID]/start/route.ts +215 -0
src/app/api/jobs/[jobID]/stop/route.ts +23 -0
src/app/api/jobs/route.ts +67 -0
src/app/api/settings/route.ts +59 -0
src/app/api/zip/route.ts +78 -0
src/app/apple-icon.png +0 -0
src/app/dashboard/page.tsx +85 -0
src/app/datasets/[datasetName]/page.tsx +190 -0
src/app/datasets/page.tsx +217 -0
src/app/favicon.ico +0 -0
src/app/globals.css +72 -0
src/app/icon.png +0 -0
src/app/icon.svg +0 -0
src/app/jobs/[jobID]/page.tsx +147 -0
src/app/jobs/new/AdvancedJob.tsx +146 -0
src/app/jobs/new/SimpleJob.tsx +973 -0
src/app/jobs/new/jobConfig.ts +167 -0
src/app/jobs/new/options.ts +441 -0
src/app/jobs/new/page.tsx +306 -0
src/app/jobs/page.tsx +49 -0
src/app/layout.tsx +50 -0
src/app/manifest.json +21 -0
src/app/page.tsx +5 -0
src/app/settings/page.tsx +264 -0
src/components/AddImagesModal.tsx +152 -0
src/components/AddSingleImageModal.tsx +141 -0
src/components/AuthWrapper.tsx +166 -0
src/components/Card.tsx +15 -0

src/app/api/auth/hf/callback/route.ts ADDED Viewed

	@@ -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 = process.env.HF_OAUTH_REDIRECT_URI || process.env.NEXT_PUBLIC_HF_OAUTH_REDIRECT_URI || `${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());
+  }
+}

src/app/api/auth/hf/login/route.ts ADDED Viewed

	@@ -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 = process.env.HF_OAUTH_REDIRECT_URI || process.env.NEXT_PUBLIC_HF_OAUTH_REDIRECT_URI || `${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;
+}

src/app/api/auth/hf/validate/route.ts ADDED Viewed

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

src/app/api/auth/route.ts ADDED Viewed

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

src/app/api/caption/get/route.ts ADDED Viewed

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

src/app/api/datasets/create/route.tsx ADDED Viewed

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

src/app/api/datasets/delete/route.tsx ADDED Viewed

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

src/app/api/datasets/list/route.ts ADDED Viewed

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

src/app/api/datasets/listImages/route.ts ADDED Viewed

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

src/app/api/datasets/upload/route.ts ADDED Viewed

	@@ -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',
+  },
+};

src/app/api/files/[...filePath]/route.ts ADDED Viewed

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

src/app/api/gpu/route.ts ADDED Viewed

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

src/app/api/hf-hub/route.ts ADDED Viewed

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

src/app/api/hf-jobs/route.ts ADDED Viewed

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

src/app/api/img/[...imagePath]/route.ts ADDED Viewed

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

src/app/api/img/caption/route.ts ADDED Viewed

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

src/app/api/img/delete/route.ts ADDED Viewed

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

src/app/api/img/upload/route.ts ADDED Viewed

	@@ -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',
+  },
+};

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

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

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

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

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

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

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

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

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

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

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

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

src/app/api/jobs/route.ts ADDED Viewed

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

src/app/api/settings/route.ts ADDED Viewed

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

src/app/api/zip/route.ts ADDED Viewed

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

src/app/apple-icon.png ADDED Viewed

src/app/dashboard/page.tsx ADDED Viewed

	@@ -0,0 +1,85 @@

+'use client';
+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" />
+      </TopBar>
+      <MainContent>
+        <div className="border border-gray-800 rounded-xl bg-gray-900 p-6 flex flex-col gap-4">
+          <div>
+            <h2 className="text-xl font-semibold text-gray-100">
+              {isAuthenticated ? `Welcome back, ${namespace || 'creator'}!` : 'Welcome to Ostris AI Toolkit'}
+            </h2>
+            <p className="text-sm text-gray-400 mt-2">
+              {isAuthenticated
+                ? 'You are signed in with Hugging Face and can manage jobs, datasets, and submissions.'
+                : 'Authenticate with Hugging Face or add a personal access token to create jobs, upload datasets, and launch training.'}
+            </p>
+          </div>
+          {isAuthenticated ? (
+            <div className="flex flex-wrap items-center gap-3 text-sm">
+              <Link
+                href="/jobs/new"
+                className="px-4 py-2 rounded-md bg-blue-600 hover:bg-blue-500 text-white transition-colors"
+              >
+                Create a Training Job
+              </Link>
+              <Link
+                href="/datasets"
+                className="px-4 py-2 rounded-md bg-gray-800 hover:bg-gray-700 text-gray-200 transition-colors"
+              >
+                Manage Datasets
+              </Link>
+              <Link
+                href="/settings"
+                className="px-4 py-2 rounded-md border border-gray-700 text-gray-300 hover:border-gray-600 transition-colors"
+              >
+                Settings
+              </Link>
+            </div>
+          ) : (
+            <div className="flex flex-wrap items-center gap-3 text-sm">
+              <HFLoginButton size="md" />
+              <Link
+                href="/settings"
+                className="text-xs text-blue-400 hover:text-blue-300"
+              >
+                Or manage tokens in Settings
+              </Link>
+            </div>
+          )}
+        </div>
+        <div className="w-full mt-6">
+          <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>
+    </>
+  );
+}

src/app/datasets/[datasetName]/page.tsx ADDED Viewed

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

src/app/datasets/page.tsx ADDED Viewed

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

src/app/favicon.ico ADDED Viewed

src/app/globals.css ADDED Viewed

	@@ -0,0 +1,72 @@

+@tailwind base;
+@tailwind components;
+@tailwind utilities;
+:root {
+  --background: #ffffff;
+  --foreground: #171717;
+}
+@media (prefers-color-scheme: dark) {
+  :root {
+    --background: #0a0a0a;
+    --foreground: #ededed;
+  }
+}
+body {
+  color: var(--foreground);
+  background: var(--background);
+  font-family: Arial, Helvetica, sans-serif;
+}
+@layer components {
+  /* control */
+  .aitk-react-select-container .aitk-react-select__control {
+    @apply flex w-full h-8 min-h-0 px-0 text-sm bg-gray-800 border border-gray-700 rounded-sm hover:border-gray-600 items-center;
+  }
+  /* selected label */
+  .aitk-react-select-container .aitk-react-select__single-value {
+    @apply flex-1 min-w-0 truncate text-sm text-neutral-200;
+  }
+  /* invisible input (keeps focus & typing, never wraps) */
+  .aitk-react-select-container .aitk-react-select__input-container {
+    @apply text-neutral-200;
+  }
+  /* focus */
+  .aitk-react-select-container .aitk-react-select__control--is-focused {
+    @apply ring-2 ring-gray-600 border-transparent hover:border-transparent shadow-none;
+  }
+  /* menu */
+  .aitk-react-select-container .aitk-react-select__menu {
+    @apply bg-gray-800 border border-gray-700;
+  }
+  /* options */
+  .aitk-react-select-container .aitk-react-select__option {
+    @apply text-sm text-neutral-200 bg-gray-800 hover:bg-gray-700;
+  }
+  /* indicator separator */
+  .aitk-react-select-container .aitk-react-select__indicator-separator {
+    @apply bg-gray-600;
+  }
+  /* indicators */
+  .aitk-react-select-container .aitk-react-select__indicators,
+  .aitk-react-select-container .aitk-react-select__indicator {
+    @apply py-0 flex items-center;
+  }
+  /* placeholder */
+  .aitk-react-select-container .aitk-react-select__placeholder {
+    @apply text-sm text-neutral-200;
+  }
+}

src/app/icon.png ADDED Viewed

src/app/icon.svg ADDED Viewed

src/app/jobs/[jobID]/page.tsx ADDED Viewed

	@@ -0,0 +1,147 @@

+'use client';
+import { useState, use } from 'react';
+import { FaChevronLeft } from 'react-icons/fa';
+import { Button } from '@headlessui/react';
+import { TopBar, MainContent } from '@/components/layout';
+import useJob from '@/hooks/useJob';
+import SampleImages, {SampleImagesMenu} from '@/components/SampleImages';
+import JobOverview from '@/components/JobOverview';
+import { redirect } from 'next/navigation';
+import { useAuth } from '@/contexts/AuthContext';
+import HFLoginButton from '@/components/HFLoginButton';
+import Link from 'next/link';
+import JobActionBar from '@/components/JobActionBar';
+import JobConfigViewer from '@/components/JobConfigViewer';
+import { JobRecord } from '@/types';
+type PageKey = 'overview' | 'samples' | 'config';
+interface Page {
+  name: string;
+  value: PageKey;
+  component: React.ComponentType<{ job: JobRecord }>;
+  menuItem?: React.ComponentType<{ job?: JobRecord | null }> | null;
+  mainCss?: string;
+}
+const pages: Page[] = [
+  {
+    name: 'Overview',
+    value: 'overview',
+    component: JobOverview,
+    mainCss: 'pt-24',
+  },
+  {
+    name: 'Samples',
+    value: 'samples',
+    component: SampleImages,
+    menuItem: SampleImagesMenu,
+    mainCss: 'pt-24',
+  },
+  {
+    name: 'Config File',
+    value: 'config',
+    component: JobConfigViewer,
+    mainCss: 'pt-[80px] px-0 pb-0',
+  },
+];
+export default function JobPage({ params }: { params: { jobID: string } }) {
+  const usableParams = use(params as any) as { jobID: string };
+  const jobID = usableParams.jobID;
+  const { job, status, refreshJob } = useJob(jobID, 5000);
+  const [pageKey, setPageKey] = useState<PageKey>('overview');
+  const { status: authStatus } = useAuth();
+  const isAuthenticated = authStatus === 'authenticated';
+  const page = pages.find(p => p.value === pageKey);
+  if (!isAuthenticated) {
+    return (
+      <>
+        <TopBar>
+          <div>
+            <Button className="text-gray-500 dark:text-gray-300 px-3 mt-1" onClick={() => redirect('/jobs')}>
+              <FaChevronLeft />
+            </Button>
+          </div>
+          <div>
+            <h1 className="text-lg">Job Details</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>Sign in with Hugging Face or add an access token to view job details.</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>
+      </>
+    );
+  }
+  return (
+    <>
+      {/* Fixed top bar */}
+      <TopBar>
+        <div>
+          <Button className="text-gray-500 dark:text-gray-300 px-3 mt-1" onClick={() => redirect('/jobs')}>
+            <FaChevronLeft />
+          </Button>
+        </div>
+        <div>
+          <h1 className="text-lg">Job: {job?.name}</h1>
+        </div>
+        <div className="flex-1"></div>
+        {job && (
+          <JobActionBar
+            job={job}
+            onRefresh={refreshJob}
+            hideView
+            afterDelete={() => {
+              redirect('/jobs');
+            }}
+          />
+        )}
+      </TopBar>
+      <MainContent className={pages.find(page => page.value === pageKey)?.mainCss}>
+        {status === 'loading' && job == null && <p>Loading...</p>}
+        {status === 'error' && job == null && <p>Error fetching job</p>}
+        {job && (
+          <>
+            {pages.map(page => {
+              const Component = page.component;
+              return page.value === pageKey ? <Component key={page.value} job={job} /> : null;
+            })}
+          </>
+        )}
+      </MainContent>
+      <div className="bg-gray-800 absolute top-12 left-0 w-full h-8 flex items-center px-2 text-sm">
+        {pages.map(page => (
+          <Button
+            key={page.value}
+            onClick={() => setPageKey(page.value)}
+            className={`px-4 py-1 h-8  ${page.value === pageKey ? 'bg-gray-300 dark:bg-gray-700' : ''}`}
+          >
+            {page.name}
+          </Button>
+        ))}
+        {
+          page?.menuItem && (
+            <>
+            <div className='flex-grow'>
+            </div>
+              <page.menuItem job={job} />
+            </>
+          )
+        }
+      </div>
+    </>
+  );
+}

src/app/jobs/new/AdvancedJob.tsx ADDED Viewed

	@@ -0,0 +1,146 @@

+'use client';
+import { useEffect, useState, useRef } from 'react';
+import { JobConfig } from '@/types';
+import YAML from 'yaml';
+import Editor, { OnMount } from '@monaco-editor/react';
+import type { editor } from 'monaco-editor';
+import { SettingsData } from '@/types';
+import { migrateJobConfig } from './jobConfig';
+type Props = {
+  jobConfig: JobConfig;
+  setJobConfig: (value: any, key?: string) => void;
+  status: 'idle' | 'saving' | 'success' | 'error';
+  handleSubmit: (event: React.FormEvent<HTMLFormElement>) => void;
+  runId: string | null;
+  gpuIDs: string | null;
+  setGpuIDs: (value: string | null) => void;
+  gpuList: any;
+  datasetOptions: any;
+  settings: SettingsData;
+};
+const isDev = process.env.NODE_ENV === 'development';
+const yamlConfig: YAML.DocumentOptions &
+  YAML.SchemaOptions &
+  YAML.ParseOptions &
+  YAML.CreateNodeOptions &
+  YAML.ToStringOptions = {
+  indent: 2,
+  lineWidth: 999999999999,
+  defaultStringType: 'QUOTE_DOUBLE',
+  defaultKeyType: 'PLAIN',
+  directives: true,
+};
+export default function AdvancedJob({ jobConfig, setJobConfig, settings }: Props) {
+  const [editorValue, setEditorValue] = useState<string>('');
+  const lastJobConfigUpdateStringRef = useRef('');
+  const editorRef = useRef<editor.IStandaloneCodeEditor | null>(null);
+  // Track if the editor has been mounted
+  const isEditorMounted = useRef(false);
+  // Handler for editor mounting
+  const handleEditorDidMount: OnMount = editor => {
+    editorRef.current = editor;
+    isEditorMounted.current = true;
+    // Initial content setup
+    try {
+      const yamlContent = YAML.stringify(jobConfig, yamlConfig);
+      setEditorValue(yamlContent);
+      lastJobConfigUpdateStringRef.current = JSON.stringify(jobConfig);
+    } catch (e) {
+      console.warn(e);
+    }
+  };
+  useEffect(() => {
+    const lastUpdate = lastJobConfigUpdateStringRef.current;
+    const currentUpdate = JSON.stringify(jobConfig);
+    // Skip if no changes or editor not yet mounted
+    if (lastUpdate === currentUpdate || !isEditorMounted.current) {
+      return;
+    }
+    try {
+      // Preserve cursor position and selection
+      const editor = editorRef.current;
+      if (editor) {
+        // Save current editor state
+        const position = editor.getPosition();
+        const selection = editor.getSelection();
+        const scrollTop = editor.getScrollTop();
+        // Update content
+        const yamlContent = YAML.stringify(jobConfig, yamlConfig);
+        // Only update if the content is actually different
+        if (yamlContent !== editor.getValue()) {
+          // Set value directly on the editor model instead of using React state
+          editor.getModel()?.setValue(yamlContent);
+          // Restore cursor position and selection
+          if (position) editor.setPosition(position);
+          if (selection) editor.setSelection(selection);
+          editor.setScrollTop(scrollTop);
+        }
+        lastJobConfigUpdateStringRef.current = currentUpdate;
+      }
+    } catch (e) {
+      console.warn(e);
+    }
+  }, [jobConfig]);
+  const handleChange = (value: string | undefined) => {
+    if (value === undefined) return;
+    try {
+      const parsed = YAML.parse(value);
+      // Don't update jobConfig if the change came from the editor itself
+      // to avoid a circular update loop
+      if (JSON.stringify(parsed) !== lastJobConfigUpdateStringRef.current) {
+        lastJobConfigUpdateStringRef.current = JSON.stringify(parsed);
+        // We have to ensure certain things are always set
+        try {
+          parsed.config.process[0].type = 'ui_trainer';
+          parsed.config.process[0].sqlite_db_path = './aitk_db.db';
+          parsed.config.process[0].training_folder = settings.TRAINING_FOLDER;
+          parsed.config.process[0].device = 'cuda';
+          parsed.config.process[0].performance_log_every = 10;
+        } catch (e) {
+          console.warn(e);
+        }
+        migrateJobConfig(parsed);
+        setJobConfig(parsed);
+      }
+    } catch (e) {
+      // Don't update on parsing errors
+      console.warn(e);
+    }
+  };
+  return (
+    <>
+      <Editor
+        height="100%"
+        width="100%"
+        defaultLanguage="yaml"
+        value={editorValue}
+        theme="vs-dark"
+        onChange={handleChange}
+        onMount={handleEditorDidMount}
+        options={{
+          minimap: { enabled: true },
+          scrollBeyondLastLine: false,
+          automaticLayout: true,
+        }}
+      />
+    </>
+  );
+}

src/app/jobs/new/SimpleJob.tsx ADDED Viewed

	@@ -0,0 +1,973 @@

+'use client';
+import { useMemo, useState } from 'react';
+import { modelArchs, ModelArch, groupedModelOptions, quantizationOptions, defaultQtype } from './options';
+import { defaultDatasetConfig } from './jobConfig';
+import { GroupedSelectOption, JobConfig, SelectOption } from '@/types';
+import { objectCopy } from '@/utils/basic';
+import { TextInput, SelectInput, Checkbox, FormGroup, NumberInput } from '@/components/formInputs';
+import Card from '@/components/Card';
+import { X } from 'lucide-react';
+import AddSingleImageModal, { openAddImageModal } from '@/components/AddSingleImageModal';
+import {FlipHorizontal2, FlipVertical2} from "lucide-react";
+import HFJobsWorkflow from '@/components/HFJobsWorkflow';
+type Props = {
+  jobConfig: JobConfig;
+  setJobConfig: (value: any, key: string) => void;
+  status: 'idle' | 'saving' | 'success' | 'error';
+  handleSubmit: (event: React.FormEvent<HTMLFormElement>) => void;
+  runId: string | null;
+  gpuIDs: string | null;
+  setGpuIDs: (value: string | null) => void;
+  gpuList: any;
+  datasetOptions: any;
+  trainingBackend?: 'local' | 'hf-jobs';
+  setTrainingBackend?: (backend: 'local' | 'hf-jobs') => void;
+  hfJobSubmitted?: boolean;
+  onHFJobComplete?: (jobId: string, localJobId?: string) => void;
+  forceHFBackend?: boolean;
+};
+const isDev = process.env.NODE_ENV === 'development';
+export default function SimpleJob({
+  jobConfig,
+  setJobConfig,
+  handleSubmit,
+  status,
+  runId,
+  gpuIDs,
+  setGpuIDs,
+  gpuList,
+  datasetOptions,
+  trainingBackend: parentTrainingBackend,
+  setTrainingBackend: parentSetTrainingBackend,
+  hfJobSubmitted,
+  onHFJobComplete,
+  forceHFBackend = false,
+}: Props) {
+  const [localTrainingBackend, setLocalTrainingBackend] = useState(forceHFBackend ? 'hf-jobs' : 'local');
+  const trainingBackend = parentTrainingBackend || localTrainingBackend;
+  const setTrainingBackend = forceHFBackend
+    ? (_: 'local' | 'hf-jobs') => undefined
+    : parentSetTrainingBackend || setLocalTrainingBackend;
+  const backendOptions = forceHFBackend
+    ? [{ value: 'hf-jobs', label: 'HF Jobs (Cloud)' }]
+    : [
+        { value: 'local', label: 'Local GPU' },
+        { value: 'hf-jobs', label: 'HF Jobs (Cloud)' },
+      ];
+  const modelArch = useMemo(() => {
+    return modelArchs.find(a => a.name === jobConfig.config.process[0].model.arch) as ModelArch;
+  }, [jobConfig.config.process[0].model.arch]);
+  const isVideoModel = !!(modelArch?.group === 'video');
+  const numTopCards = useMemo(() => {
+    let count = 4; // job settings, model config, target config, save config
+    if (modelArch?.additionalSections?.includes('model.multistage')) {
+      count += 1; // add multistage card
+    }
+    if (!modelArch?.disableSections?.includes('model.quantize')) {
+      count += 1; // add quantization card
+    }
+    return count;
+  }, [modelArch]);
+  let topBarClass = 'grid grid-cols-1 md:grid-cols-2 lg:grid-cols-4 xl:grid-cols-4 gap-6';
+  if (numTopCards == 5) {
+    topBarClass = 'grid grid-cols-1 md:grid-cols-2 lg:grid-cols-3 xl:grid-cols-5 gap-6';
+  }
+  if (numTopCards == 6) {
+    topBarClass = 'grid grid-cols-1 md:grid-cols-2 lg:grid-cols-3 xl:grid-cols-3 2xl:grid-cols-6 gap-6';
+  }
+  const transformerQuantizationOptions: GroupedSelectOption[] | SelectOption[] = useMemo(() => {
+    const hasARA = modelArch?.accuracyRecoveryAdapters && Object.keys(modelArch.accuracyRecoveryAdapters).length > 0;
+    if (!hasARA) {
+      return quantizationOptions;
+    }
+    let newQuantizationOptions = [
+      {
+        label: 'Standard',
+        options: [quantizationOptions[0], quantizationOptions[1]],
+      },
+    ];
+    // add ARAs if they exist for the model
+    let ARAs: SelectOption[] = [];
+    if (modelArch.accuracyRecoveryAdapters) {
+      for (const [label, value] of Object.entries(modelArch.accuracyRecoveryAdapters)) {
+         ARAs.push({ value, label });
+      }
+    }
+    if (ARAs.length > 0) {
+      newQuantizationOptions.push({
+        label: 'Accuracy Recovery Adapters',
+        options: ARAs,
+      });
+    }
+    let additionalQuantizationOptions: SelectOption[] = [];
+    // add the quantization options if they are not already included
+    for (let i = 2; i < quantizationOptions.length; i++) {
+      const option = quantizationOptions[i];
+      additionalQuantizationOptions.push(option);
+    }
+    if (additionalQuantizationOptions.length > 0) {
+      newQuantizationOptions.push({
+        label: 'Additional Quantization Options',
+        options: additionalQuantizationOptions,
+      });
+    }
+    return newQuantizationOptions;
+  }, [modelArch]);
+  return (
+    <>
+      <form onSubmit={handleSubmit} className="space-y-8">
+        <div className={topBarClass}>
+          <Card title="Job">
+            <TextInput
+              label="Training Name"
+              value={jobConfig.config.name}
+              docKey="config.name"
+              onChange={value => setJobConfig(value, 'config.name')}
+              placeholder="Enter training name"
+              disabled={runId !== null}
+              required
+            />
+            <SelectInput
+              label="Training Backend"
+              value={trainingBackend}
+              onChange={(value) => {
+                setTrainingBackend(value);
+              }}
+              options={backendOptions}
+              disabled={forceHFBackend}
+            />
+            {trainingBackend === 'local' && (
+              <SelectInput
+                label="GPU ID"
+                value={`${gpuIDs}`}
+                docKey="gpuids"
+                onChange={value => setGpuIDs(value)}
+                options={gpuList.map((gpu: any) => ({ value: `${gpu.index}`, label: `GPU #${gpu.index}` }))}
+              />
+            )}
+            <TextInput
+              label="Trigger Word"
+              value={jobConfig.config.process[0].trigger_word || ''}
+              docKey="config.process[0].trigger_word"
+              onChange={(value: string | null) => {
+                if (value?.trim() === '') {
+                  value = null;
+                }
+                setJobConfig(value, 'config.process[0].trigger_word');
+              }}
+              placeholder=""
+              required
+            />
+            {trainingBackend === 'hf-jobs' && (
+              <div className={`mt-4 p-3 rounded ${
+                hfJobSubmitted
+                  ? 'bg-green-900/20 border border-green-700'
+                  : 'bg-yellow-900/20 border border-yellow-700'
+              }`}>
+                <p className={`text-sm ${
+                  hfJobSubmitted ? 'text-green-400' : 'text-yellow-400'
+                }`}>
+                  {hfJobSubmitted
+                    ? '✓ HF Job already submitted! You can modify settings and resubmit if needed.'
+                    : '⏳ HF Job ready for submission. Submit to the cloud below.'
+                  }
+                </p>
+              </div>
+            )}
+          </Card>
+          {/* Model Configuration Section */}
+          <Card title="Model">
+            <SelectInput
+              label="Model Architecture"
+              value={jobConfig.config.process[0].model.arch}
+              onChange={value => {
+                const currentArch = modelArchs.find(a => a.name === jobConfig.config.process[0].model.arch);
+                if (!currentArch || currentArch.name === value) {
+                  return;
+                }
+                // update the defaults when a model is selected
+                const newArch = modelArchs.find(model => model.name === value);
+                // update vram setting
+                if (!newArch?.additionalSections?.includes('model.low_vram')) {
+                  setJobConfig(false, 'config.process[0].model.low_vram');
+                }
+                // revert defaults from previous model
+                for (const key in currentArch.defaults) {
+                  setJobConfig(currentArch.defaults[key][1], key);
+                }
+                if (newArch?.defaults) {
+                  for (const key in newArch.defaults) {
+                    setJobConfig(newArch.defaults[key][0], key);
+                  }
+                }
+                // set new model
+                setJobConfig(value, 'config.process[0].model.arch');
+                // update datasets
+                const hasControlPath = newArch?.additionalSections?.includes('datasets.control_path') || false;
+                const hasNumFrames = newArch?.additionalSections?.includes('datasets.num_frames') || false;
+                const controls = newArch?.controls ?? [];
+                const datasets = jobConfig.config.process[0].datasets.map(dataset => {
+                  const newDataset = objectCopy(dataset);
+                  newDataset.controls = controls;
+                  if (!hasControlPath) {
+                    newDataset.control_path = null; // reset control path if not applicable
+                  }
+                  if (!hasNumFrames) {
+                    newDataset.num_frames = 1; // reset num_frames if not applicable
+                  }
+                  return newDataset;
+                });
+                setJobConfig(datasets, 'config.process[0].datasets');
+                // update samples
+                const hasSampleCtrlImg = newArch?.additionalSections?.includes('sample.ctrl_img') || false;
+                const samples = jobConfig.config.process[0].sample.samples.map(sample => {
+                  const newSample = objectCopy(sample);
+                  if (!hasSampleCtrlImg) {
+                    delete newSample.ctrl_img; // remove ctrl_img if not applicable
+                  }
+                  return newSample;
+                });
+                setJobConfig(samples, 'config.process[0].sample.samples');
+              }}
+              options={groupedModelOptions}
+            />
+            <TextInput
+              label="Name or Path"
+              value={jobConfig.config.process[0].model.name_or_path}
+              docKey="config.process[0].model.name_or_path"
+              onChange={(value: string | null) => {
+                if (value?.trim() === '') {
+                  value = null;
+                }
+                setJobConfig(value, 'config.process[0].model.name_or_path');
+              }}
+              placeholder=""
+              required
+            />
+            {modelArch?.additionalSections?.includes('model.low_vram') && (
+              <FormGroup label="Options">
+                <Checkbox
+                  label="Low VRAM"
+                  checked={jobConfig.config.process[0].model.low_vram}
+                  onChange={value => setJobConfig(value, 'config.process[0].model.low_vram')}
+                />
+              </FormGroup>
+            )}
+          </Card>
+          {modelArch?.disableSections?.includes('model.quantize') ? null : (
+            <Card title="Quantization">
+              <SelectInput
+                label="Transformer"
+                value={jobConfig.config.process[0].model.quantize ? jobConfig.config.process[0].model.qtype : ''}
+                onChange={value => {
+                  if (value === '') {
+                    setJobConfig(false, 'config.process[0].model.quantize');
+                    value = defaultQtype;
+                  } else {
+                    setJobConfig(true, 'config.process[0].model.quantize');
+                  }
+                  setJobConfig(value, 'config.process[0].model.qtype');
+                }}
+                options={transformerQuantizationOptions}
+              />
+              <SelectInput
+                label="Text Encoder"
+                value={jobConfig.config.process[0].model.quantize_te ? jobConfig.config.process[0].model.qtype_te : ''}
+                onChange={value => {
+                  if (value === '') {
+                    setJobConfig(false, 'config.process[0].model.quantize_te');
+                    value = defaultQtype;
+                  } else {
+                    setJobConfig(true, 'config.process[0].model.quantize_te');
+                  }
+                  setJobConfig(value, 'config.process[0].model.qtype_te');
+                }}
+                options={quantizationOptions}
+              />
+            </Card>
+          )}
+          {modelArch?.additionalSections?.includes('model.multistage') && (
+            <Card title="Multistage">
+              <FormGroup label="Stages to Train" docKey={'model.multistage'}>
+                <Checkbox
+                  label="High Noise"
+                  checked={jobConfig.config.process[0].model.model_kwargs?.train_high_noise || false}
+                  onChange={value => setJobConfig(value, 'config.process[0].model.model_kwargs.train_high_noise')}
+                />
+                <Checkbox
+                  label="Low Noise"
+                  checked={jobConfig.config.process[0].model.model_kwargs?.train_low_noise || false}
+                  onChange={value => setJobConfig(value, 'config.process[0].model.model_kwargs.train_low_noise')}
+                />
+              </FormGroup>
+              <NumberInput
+                  label="Switch Every"
+                  value={jobConfig.config.process[0].train.switch_boundary_every}
+                  onChange={value => setJobConfig(value, 'config.process[0].train.switch_boundary_every')}
+                  placeholder="eg. 1"
+                  docKey={'train.switch_boundary_every'}
+                  min={1}
+                  required
+                />
+            </Card>
+          )}
+          <Card title="Target">
+            <SelectInput
+              label="Target Type"
+              value={jobConfig.config.process[0].network?.type ?? 'lora'}
+              onChange={value => setJobConfig(value, 'config.process[0].network.type')}
+              options={[
+                { value: 'lora', label: 'LoRA' },
+                { value: 'lokr', label: 'LoKr' },
+              ]}
+            />
+            {jobConfig.config.process[0].network?.type == 'lokr' && (
+              <SelectInput
+                label="LoKr Factor"
+                value={`${jobConfig.config.process[0].network?.lokr_factor ?? -1}`}
+                onChange={value => setJobConfig(parseInt(value), 'config.process[0].network.lokr_factor')}
+                options={[
+                  { value: '-1', label: 'Auto' },
+                  { value: '4', label: '4' },
+                  { value: '8', label: '8' },
+                  { value: '16', label: '16' },
+                  { value: '32', label: '32' },
+                ]}
+              />
+            )}
+            {jobConfig.config.process[0].network?.type == 'lora' && (
+              <>
+                <NumberInput
+                  label="Linear Rank"
+                  value={jobConfig.config.process[0].network.linear}
+                  onChange={value => {
+                    console.log('onChange', value);
+                    setJobConfig(value, 'config.process[0].network.linear');
+                    setJobConfig(value, 'config.process[0].network.linear_alpha');
+                  }}
+                  placeholder="eg. 16"
+                  min={0}
+                  max={1024}
+                  required
+                />
+                {modelArch?.disableSections?.includes('network.conv') ? null : (
+                  <NumberInput
+                    label="Conv Rank"
+                    value={jobConfig.config.process[0].network.conv}
+                    onChange={value => {
+                      console.log('onChange', value);
+                      setJobConfig(value, 'config.process[0].network.conv');
+                      setJobConfig(value, 'config.process[0].network.conv_alpha');
+                    }}
+                    placeholder="eg. 16"
+                    min={0}
+                    max={1024}
+                  />
+                )}
+              </>
+            )}
+          </Card>
+          <Card title="Save">
+            <SelectInput
+              label="Data Type"
+              value={jobConfig.config.process[0].save.dtype}
+              onChange={value => setJobConfig(value, 'config.process[0].save.dtype')}
+              options={[
+                { value: 'bf16', label: 'BF16' },
+                { value: 'fp16', label: 'FP16' },
+                { value: 'fp32', label: 'FP32' },
+              ]}
+            />
+            <NumberInput
+              label="Save Every"
+              value={jobConfig.config.process[0].save.save_every}
+              onChange={value => setJobConfig(value, 'config.process[0].save.save_every')}
+              placeholder="eg. 250"
+              min={1}
+              required
+            />
+            <NumberInput
+              label="Max Step Saves to Keep"
+              value={jobConfig.config.process[0].save.max_step_saves_to_keep}
+              onChange={value => setJobConfig(value, 'config.process[0].save.max_step_saves_to_keep')}
+              placeholder="eg. 4"
+              min={1}
+              required
+            />
+          </Card>
+        </div>
+        <div>
+          <Card title="Training">
+            <div className="grid grid-cols-1 md:grid-cols-3 lg:grid-cols-5 gap-6">
+              <div>
+                <NumberInput
+                  label="Batch Size"
+                  value={jobConfig.config.process[0].train.batch_size}
+                  onChange={value => setJobConfig(value, 'config.process[0].train.batch_size')}
+                  placeholder="eg. 4"
+                  min={1}
+                  required
+                />
+                <NumberInput
+                  label="Gradient Accumulation"
+                  className="pt-2"
+                  value={jobConfig.config.process[0].train.gradient_accumulation}
+                  onChange={value => setJobConfig(value, 'config.process[0].train.gradient_accumulation')}
+                  placeholder="eg. 1"
+                  min={1}
+                  required
+                />
+                <NumberInput
+                  label="Steps"
+                  className="pt-2"
+                  value={jobConfig.config.process[0].train.steps}
+                  onChange={value => setJobConfig(value, 'config.process[0].train.steps')}
+                  placeholder="eg. 2000"
+                  min={1}
+                  required
+                />
+              </div>
+              <div>
+                <SelectInput
+                  label="Optimizer"
+                  value={jobConfig.config.process[0].train.optimizer}
+                  onChange={value => setJobConfig(value, 'config.process[0].train.optimizer')}
+                  options={[
+                    { value: 'adamw8bit', label: 'AdamW8Bit' },
+                    { value: 'adafactor', label: 'Adafactor' },
+                  ]}
+                />
+                <NumberInput
+                  label="Learning Rate"
+                  className="pt-2"
+                  value={jobConfig.config.process[0].train.lr}
+                  onChange={value => setJobConfig(value, 'config.process[0].train.lr')}
+                  placeholder="eg. 0.0001"
+                  min={0}
+                  required
+                />
+                <NumberInput
+                  label="Weight Decay"
+                  className="pt-2"
+                  value={jobConfig.config.process[0].train.optimizer_params.weight_decay}
+                  onChange={value => setJobConfig(value, 'config.process[0].train.optimizer_params.weight_decay')}
+                  placeholder="eg. 0.0001"
+                  min={0}
+                  required
+                />
+              </div>
+              <div>
+                {modelArch?.disableSections?.includes('train.timestep_type') ? null : (
+                  <SelectInput
+                    label="Timestep Type"
+                    value={jobConfig.config.process[0].train.timestep_type}
+                    disabled={modelArch?.disableSections?.includes('train.timestep_type') || false}
+                    onChange={value => setJobConfig(value, 'config.process[0].train.timestep_type')}
+                    options={[
+                      { value: 'sigmoid', label: 'Sigmoid' },
+                      { value: 'linear', label: 'Linear' },
+                      { value: 'shift', label: 'Shift' },
+                      { value: 'weighted', label: 'Weighted' },
+                    ]}
+                  />
+                )}
+                <SelectInput
+                  label="Timestep Bias"
+                  className="pt-2"
+                  value={jobConfig.config.process[0].train.content_or_style}
+                  onChange={value => setJobConfig(value, 'config.process[0].train.content_or_style')}
+                  options={[
+                    { value: 'balanced', label: 'Balanced' },
+                    { value: 'content', label: 'High Noise' },
+                    { value: 'style', label: 'Low Noise' },
+                  ]}
+                />
+                <SelectInput
+                  label="Noise Scheduler"
+                  className="pt-2"
+                  value={jobConfig.config.process[0].train.noise_scheduler}
+                  onChange={value => setJobConfig(value, 'config.process[0].train.noise_scheduler')}
+                  options={[
+                    { value: 'flowmatch', label: 'FlowMatch' },
+                    { value: 'ddpm', label: 'DDPM' },
+                  ]}
+                />
+              </div>
+              <div>
+                <FormGroup label="EMA (Exponential Moving Average)">
+                  <Checkbox
+                    label="Use EMA"
+                    className="pt-1"
+                    checked={jobConfig.config.process[0].train.ema_config?.use_ema || false}
+                    onChange={value => setJobConfig(value, 'config.process[0].train.ema_config.use_ema')}
+                  />
+                </FormGroup>
+                {jobConfig.config.process[0].train.ema_config?.use_ema && (
+                  <NumberInput
+                    label="EMA Decay"
+                    className="pt-2"
+                    value={jobConfig.config.process[0].train.ema_config?.ema_decay as number}
+                    onChange={value => setJobConfig(value, 'config.process[0].train.ema_config?.ema_decay')}
+                    placeholder="eg. 0.99"
+                    min={0}
+                  />
+                )}
+                <FormGroup label="Text Encoder Optimizations" className="pt-2">
+                  <Checkbox
+                    label="Unload TE"
+                    checked={jobConfig.config.process[0].train.unload_text_encoder || false}
+                    docKey={'train.unload_text_encoder'}
+                    onChange={value => {
+                      setJobConfig(value, 'config.process[0].train.unload_text_encoder');
+                      if (value) {
+                        setJobConfig(false, 'config.process[0].train.cache_text_embeddings');
+                      }
+                    }}
+                  />
+                  <Checkbox
+                    label="Cache Text Embeddings"
+                    checked={jobConfig.config.process[0].train.cache_text_embeddings || false}
+                    docKey={'train.cache_text_embeddings'}
+                    onChange={value => {
+                      setJobConfig(value, 'config.process[0].train.cache_text_embeddings');
+                      if (value) {
+                        setJobConfig(false, 'config.process[0].train.unload_text_encoder');
+                      }
+                    }}
+                  />
+                </FormGroup>
+              </div>
+              <div>
+                <FormGroup label="Regularization">
+                  <Checkbox
+                    label="Differtial Output Preservation"
+                    className="pt-1"
+                    checked={jobConfig.config.process[0].train.diff_output_preservation || false}
+                    onChange={value => setJobConfig(value, 'config.process[0].train.diff_output_preservation')}
+                  />
+                </FormGroup>
+                {jobConfig.config.process[0].train.diff_output_preservation && (
+                  <>
+                    <NumberInput
+                      label="DOP Loss Multiplier"
+                      className="pt-2"
+                      value={jobConfig.config.process[0].train.diff_output_preservation_multiplier as number}
+                      onChange={value =>
+                        setJobConfig(value, 'config.process[0].train.diff_output_preservation_multiplier')
+                      }
+                      placeholder="eg. 1.0"
+                      min={0}
+                    />
+                    <TextInput
+                      label="DOP Preservation Class"
+                      className="pt-2"
+                      value={jobConfig.config.process[0].train.diff_output_preservation_class as string}
+                      onChange={value => setJobConfig(value, 'config.process[0].train.diff_output_preservation_class')}
+                      placeholder="eg. woman"
+                    />
+                  </>
+                )}
+              </div>
+            </div>
+          </Card>
+        </div>
+        <div>
+          <Card title="Datasets">
+            <>
+              {jobConfig.config.process[0].datasets.map((dataset, i) => (
+                <div key={i} className="p-4 rounded-lg bg-gray-800 relative">
+                  <button
+                    type="button"
+                    onClick={() =>
+                      setJobConfig(
+                        jobConfig.config.process[0].datasets.filter((_, index) => index !== i),
+                        'config.process[0].datasets',
+                      )
+                    }
+                    className="absolute top-2 right-2 bg-red-800 hover:bg-red-700 rounded-full p-1 text-sm transition-colors"
+                  >
+                    <X />
+                  </button>
+                  <h2 className="text-lg font-bold mb-4">Dataset {i + 1}</h2>
+                  <div className="grid grid-cols-1 md:grid-cols-2 lg:grid-cols-4 gap-6">
+                    <div>
+                      <SelectInput
+                        label="Dataset"
+                        value={dataset.folder_path}
+                        onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].folder_path`)}
+                        options={datasetOptions}
+                      />
+                      {modelArch?.additionalSections?.includes('datasets.control_path') && (
+                        <SelectInput
+                          label="Control Dataset"
+                          docKey="datasets.control_path"
+                          value={dataset.control_path ?? ''}
+                          className="pt-2"
+                          onChange={value =>
+                            setJobConfig(value == '' ? null : value, `config.process[0].datasets[${i}].control_path`)
+                          }
+                          options={[{ value: '', label: <>&nbsp;</> }, ...datasetOptions]}
+                        />
+                      )}
+                      <NumberInput
+                        label="LoRA Weight"
+                        value={dataset.network_weight}
+                        className="pt-2"
+                        onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].network_weight`)}
+                        placeholder="eg. 1.0"
+                      />
+                    </div>
+                    <div>
+                      <TextInput
+                        label="Default Caption"
+                        value={dataset.default_caption}
+                        onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].default_caption`)}
+                        placeholder="eg. A photo of a cat"
+                      />
+                      <NumberInput
+                        label="Caption Dropout Rate"
+                        className="pt-2"
+                        value={dataset.caption_dropout_rate}
+                        onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].caption_dropout_rate`)}
+                        placeholder="eg. 0.05"
+                        min={0}
+                        required
+                      />
+                      {modelArch?.additionalSections?.includes('datasets.num_frames') && (
+                        <NumberInput
+                          label="Num Frames"
+                          className="pt-2"
+                          docKey="datasets.num_frames"
+                          value={dataset.num_frames}
+                          onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].num_frames`)}
+                          placeholder="eg. 41"
+                          min={1}
+                          required
+                        />
+                      )}
+                    </div>
+                    <div>
+                      <FormGroup label="Settings" className="">
+                        <Checkbox
+                          label="Cache Latents"
+                          checked={dataset.cache_latents_to_disk || false}
+                          onChange={value =>
+                            setJobConfig(value, `config.process[0].datasets[${i}].cache_latents_to_disk`)
+                          }
+                        />
+                        <Checkbox
+                          label="Is Regularization"
+                          checked={dataset.is_reg || false}
+                          onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].is_reg`)}
+                        />
+                        {modelArch?.additionalSections?.includes('datasets.do_i2v') && (
+                          <Checkbox
+                            label="Do I2V"
+                            checked={dataset.do_i2v || false}
+                            onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].do_i2v`)}
+                            docKey="datasets.do_i2v"
+                          />
+                        )}
+                      </FormGroup>
+                      <FormGroup label="Flipping" docKey={'datasets.flip'} className="mt-2">
+                        <Checkbox
+                          label={<>Flip X <FlipHorizontal2 className="inline-block w-4 h-4 ml-1" /></>}
+                          checked={dataset.flip_x || false}
+                          onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].flip_x`)}
+                        />
+                        <Checkbox
+                          label={<>Flip Y <FlipVertical2 className="inline-block w-4 h-4 ml-1" /></>}
+                          checked={dataset.flip_y || false}
+                          onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].flip_y`)}
+                        />
+                      </FormGroup>
+                    </div>
+                    <div>
+                      <FormGroup label="Resolutions" className="pt-2">
+                        <div className="grid grid-cols-2 gap-2">
+                          {[
+                            [256, 512, 768],
+                            [1024, 1280, 1536],
+                          ].map(resGroup => (
+                            <div key={resGroup[0]} className="space-y-2">
+                              {resGroup.map(res => (
+                                <Checkbox
+                                  key={res}
+                                  label={res.toString()}
+                                  checked={dataset.resolution.includes(res)}
+                                  onChange={value => {
+                                    const resolutions = dataset.resolution.includes(res)
+                                      ? dataset.resolution.filter(r => r !== res)
+                                      : [...dataset.resolution, res];
+                                    setJobConfig(resolutions, `config.process[0].datasets[${i}].resolution`);
+                                  }}
+                                />
+                              ))}
+                            </div>
+                          ))}
+                        </div>
+                      </FormGroup>
+                    </div>
+                  </div>
+                </div>
+              ))}
+              <button
+                type="button"
+                onClick={() => {
+                  const newDataset = objectCopy(defaultDatasetConfig);
+                  // automaticallt add the controls for a new dataset
+                  const controls = modelArch?.controls ?? [];
+                  newDataset.controls = controls;
+                  setJobConfig([...jobConfig.config.process[0].datasets, newDataset], 'config.process[0].datasets');
+                }}
+                className="w-full px-4 py-2 bg-gray-700 hover:bg-gray-600 rounded-lg transition-colors"
+              >
+                Add Dataset
+              </button>
+            </>
+          </Card>
+        </div>
+        <div>
+          <Card title="Sample">
+            <div
+              className={
+                isVideoModel
+                  ? 'grid grid-cols-1 md:grid-cols-3 lg:grid-cols-5 gap-6'
+                  : 'grid grid-cols-1 md:grid-cols-2 lg:grid-cols-4 gap-6'
+              }
+            >
+              <div>
+                <NumberInput
+                  label="Sample Every"
+                  value={jobConfig.config.process[0].sample.sample_every}
+                  onChange={value => setJobConfig(value, 'config.process[0].sample.sample_every')}
+                  placeholder="eg. 250"
+                  min={1}
+                  required
+                />
+                <SelectInput
+                  label="Sampler"
+                  className="pt-2"
+                  value={jobConfig.config.process[0].sample.sampler}
+                  onChange={value => setJobConfig(value, 'config.process[0].sample.sampler')}
+                  options={[
+                    { value: 'flowmatch', label: 'FlowMatch' },
+                    { value: 'ddpm', label: 'DDPM' },
+                  ]}
+                />
+                <NumberInput
+                  label="Guidance Scale"
+                  value={jobConfig.config.process[0].sample.guidance_scale}
+                  onChange={value => setJobConfig(value, 'config.process[0].sample.guidance_scale')}
+                  placeholder="eg. 1.0"
+                  className="pt-2"
+                  min={0}
+                  required
+                />
+                <NumberInput
+                  label="Sample Steps"
+                  value={jobConfig.config.process[0].sample.sample_steps}
+                  onChange={value => setJobConfig(value, 'config.process[0].sample.sample_steps')}
+                  placeholder="eg. 1"
+                  className="pt-2"
+                  min={1}
+                  required
+                />
+              </div>
+              <div>
+                <NumberInput
+                  label="Width"
+                  value={jobConfig.config.process[0].sample.width}
+                  onChange={value => setJobConfig(value, 'config.process[0].sample.width')}
+                  placeholder="eg. 1024"
+                  min={0}
+                  required
+                />
+                <NumberInput
+                  label="Height"
+                  value={jobConfig.config.process[0].sample.height}
+                  onChange={value => setJobConfig(value, 'config.process[0].sample.height')}
+                  placeholder="eg. 1024"
+                  className="pt-2"
+                  min={0}
+                  required
+                />
+                {isVideoModel && (
+                  <div>
+                    <NumberInput
+                      label="Num Frames"
+                      value={jobConfig.config.process[0].sample.num_frames}
+                      onChange={value => setJobConfig(value, 'config.process[0].sample.num_frames')}
+                      placeholder="eg. 0"
+                      className="pt-2"
+                      min={0}
+                      required
+                    />
+                    <NumberInput
+                      label="FPS"
+                      value={jobConfig.config.process[0].sample.fps}
+                      onChange={value => setJobConfig(value, 'config.process[0].sample.fps')}
+                      placeholder="eg. 0"
+                      className="pt-2"
+                      min={0}
+                      required
+                    />
+                  </div>
+                )}
+              </div>
+              <div>
+                <NumberInput
+                  label="Seed"
+                  value={jobConfig.config.process[0].sample.seed}
+                  onChange={value => setJobConfig(value, 'config.process[0].sample.seed')}
+                  placeholder="eg. 0"
+                  min={0}
+                  required
+                />
+                <Checkbox
+                  label="Walk Seed"
+                  className="pt-4 pl-2"
+                  checked={jobConfig.config.process[0].sample.walk_seed}
+                  onChange={value => setJobConfig(value, 'config.process[0].sample.walk_seed')}
+                />
+              </div>
+              <div>
+                <FormGroup label="Advanced Sampling" className="pt-2">
+                  <div>
+                    <Checkbox
+                      label="Skip First Sample"
+                      className="pt-4"
+                      checked={jobConfig.config.process[0].train.skip_first_sample || false}
+                      onChange={value => setJobConfig(value, 'config.process[0].train.skip_first_sample')}
+                    />
+                  </div>
+                  <div>
+                    <Checkbox
+                      label="Disable Sampling"
+                      className="pt-1"
+                      checked={jobConfig.config.process[0].train.disable_sampling || false}
+                      onChange={value => setJobConfig(value, 'config.process[0].train.disable_sampling')}
+                    />
+                  </div>
+                </FormGroup>
+              </div>
+            </div>
+            <FormGroup label={`Sample Prompts (${jobConfig.config.process[0].sample.samples.length})`} className="pt-2">
+              <div></div>
+            </FormGroup>
+            {jobConfig.config.process[0].sample.samples.map((sample, i) => (
+              <div key={i} className="rounded-lg pl-4 pr-1 mb-4 bg-gray-950">
+                <div className="flex items-center space-x-2">
+                  <div className="flex-1">
+                    <div className="flex">
+                      <div className="flex-1">
+                        <TextInput
+                          label={`Prompt`}
+                          value={sample.prompt}
+                          onChange={value => setJobConfig(value, `config.process[0].sample.samples[${i}].prompt`)}
+                          placeholder="Enter prompt"
+                          required
+                        />
+                      </div>
+                      {modelArch?.additionalSections?.includes('sample.ctrl_img') && (
+                        <div
+                          className="h-14 w-14 mt-2 ml-4 border border-gray-500 flex items-center justify-center rounded cursor-pointer hover:bg-gray-700 transition-colors"
+                          style={{
+                            backgroundImage: sample.ctrl_img
+                              ? `url(${`/api/img/${encodeURIComponent(sample.ctrl_img)}`})`
+                              : 'none',
+                            backgroundSize: 'cover',
+                            backgroundPosition: 'center',
+                            marginBottom: '-1rem',
+                          }}
+                          onClick={() => {
+                            openAddImageModal(imagePath => {
+                              console.log('Selected image path:', imagePath);
+                              if (!imagePath) return;
+                              setJobConfig(imagePath, `config.process[0].sample.samples[${i}].ctrl_img`);
+                            });
+                          }}
+                        >
+                          {!sample.ctrl_img && (
+                            <div className="text-gray-400 text-xs text-center font-bold">Add Control Image</div>
+                          )}
+                        </div>
+                      )}
+                    </div>
+                    <div className="pb-4"></div>
+                  </div>
+                  <div>
+                    <button
+                      type="button"
+                      onClick={() =>
+                        setJobConfig(
+                          jobConfig.config.process[0].sample.samples.filter((_, index) => index !== i),
+                          'config.process[0].sample.samples',
+                        )
+                      }
+                      className="rounded-full p-1 text-sm"
+                    >
+                      <X />
+                    </button>
+                  </div>
+                </div>
+              </div>
+            ))}
+            <button
+              type="button"
+              onClick={() =>
+                setJobConfig(
+                  [...jobConfig.config.process[0].sample.samples, { prompt: '' }],
+                  'config.process[0].sample.samples',
+                )
+              }
+              className="w-full px-4 py-2 bg-gray-700 hover:bg-gray-600 rounded-lg transition-colors"
+            >
+              Add Prompt
+            </button>
+          </Card>
+        </div>
+        {status === 'success' && <p className="text-green-500 text-center">Training saved successfully!</p>}
+        {status === 'error' && <p className="text-red-500 text-center">Error saving training. Please try again.</p>}
+      </form>
+      {trainingBackend === 'hf-jobs' && (
+        <div className="mt-8">
+          <HFJobsWorkflow
+            jobConfig={jobConfig}
+            onComplete={(jobId, localJobId) => {
+              console.log('HF Job submitted:', jobId, 'Local job ID:', localJobId);
+              if (onHFJobComplete) {
+                onHFJobComplete(jobId, localJobId);
+              }
+            }}
+          />
+        </div>
+      )}
+      <AddSingleImageModal />
+    </>
+  );
+}

src/app/jobs/new/jobConfig.ts ADDED Viewed

	@@ -0,0 +1,167 @@

+import { JobConfig, DatasetConfig } from '@/types';
+export const defaultDatasetConfig: DatasetConfig = {
+  folder_path: '/path/to/images/folder',
+  control_path: null,
+  mask_path: null,
+  mask_min_value: 0.1,
+  default_caption: '',
+  caption_ext: 'txt',
+  caption_dropout_rate: 0.05,
+  cache_latents_to_disk: false,
+  is_reg: false,
+  network_weight: 1,
+  resolution: [512, 768, 1024],
+  controls: [],
+  shrink_video_to_frames: true,
+  num_frames: 1,
+  do_i2v: true,
+  flip_x: false,
+  flip_y: false,
+};
+export const defaultJobConfig: JobConfig = {
+  job: 'extension',
+  config: {
+    name: 'my_first_lora_v1',
+    process: [
+      {
+        type: 'ui_trainer',
+        training_folder: 'output',
+        sqlite_db_path: './aitk_db.db',
+        device: 'cuda',
+        trigger_word: null,
+        performance_log_every: 10,
+        network: {
+          type: 'lora',
+          linear: 32,
+          linear_alpha: 32,
+          conv: 16,
+          conv_alpha: 16,
+          lokr_full_rank: true,
+          lokr_factor: -1,
+          network_kwargs: {
+            ignore_if_contains: [],
+          },
+        },
+        save: {
+          dtype: 'bf16',
+          save_every: 250,
+          max_step_saves_to_keep: 4,
+          save_format: 'diffusers',
+          push_to_hub: false,
+        },
+        datasets: [defaultDatasetConfig],
+        train: {
+          batch_size: 1,
+          bypass_guidance_embedding: true,
+          steps: 3000,
+          gradient_accumulation: 1,
+          train_unet: true,
+          train_text_encoder: false,
+          gradient_checkpointing: true,
+          noise_scheduler: 'flowmatch',
+          optimizer: 'adamw8bit',
+          timestep_type: 'sigmoid',
+          content_or_style: 'balanced',
+          optimizer_params: {
+            weight_decay: 1e-4,
+          },
+          unload_text_encoder: false,
+          cache_text_embeddings: false,
+          lr: 0.0001,
+          ema_config: {
+            use_ema: false,
+            ema_decay: 0.99,
+          },
+          skip_first_sample: false,
+          disable_sampling: false,
+          dtype: 'bf16',
+          diff_output_preservation: false,
+          diff_output_preservation_multiplier: 1.0,
+          diff_output_preservation_class: 'person',
+          switch_boundary_every: 1,
+        },
+        model: {
+          name_or_path: 'ostris/Flex.1-alpha',
+          quantize: true,
+          qtype: 'qfloat8',
+          quantize_te: true,
+          qtype_te: 'qfloat8',
+          arch: 'flex1',
+          low_vram: false,
+          model_kwargs: {},
+        },
+        sample: {
+          sampler: 'flowmatch',
+          sample_every: 250,
+          width: 1024,
+          height: 1024,
+          samples: [
+            {
+              prompt: 'woman with red hair, playing chess at the park, bomb going off in the background'
+            },
+            {
+              prompt: 'a woman holding a coffee cup, in a beanie, sitting at a cafe',
+            },
+            {
+              prompt: 'a horse is a DJ at a night club, fish eye lens, smoke machine, lazer lights, holding a martini',
+            },
+            {
+              prompt: 'a man showing off his cool new t shirt at the beach, a shark is jumping out of the water in the background',
+            },
+            {
+              prompt: 'a bear building a log cabin in the snow covered mountains',
+            },
+            {
+              prompt: 'woman playing the guitar, on stage, singing a song, laser lights, punk rocker',
+            },
+            {
+              prompt: 'hipster man with a beard, building a chair, in a wood shop',
+            },
+            {
+              prompt: 'photo of a man, white background, medium shot, modeling clothing, studio lighting, white backdrop',
+            },
+            {
+              prompt: "a man holding a sign that says, 'this is a sign'",
+            },
+            {
+              prompt: 'a bulldog, in a post apocalyptic world, with a shotgun, in a leather jacket, in a desert, with a motorcycle',
+            },
+          ],
+          neg: '',
+          seed: 42,
+          walk_seed: true,
+          guidance_scale: 4,
+          sample_steps: 25,
+          num_frames: 1,
+          fps: 1,
+        },
+      },
+    ],
+  },
+  meta: {
+    name: '[name]',
+    version: '1.0',
+  },
+};
+export const migrateJobConfig = (jobConfig: JobConfig): JobConfig => {
+  // upgrade prompt strings to samples
+  if (
+    jobConfig?.config?.process &&
+    jobConfig.config.process[0]?.sample &&
+    Array.isArray(jobConfig.config.process[0].sample.prompts) &&
+    jobConfig.config.process[0].sample.prompts.length > 0
+  ) {
+    let newSamples = [];
+    for (const prompt of jobConfig.config.process[0].sample.prompts) {
+      newSamples.push({
+        prompt: prompt,
+      });
+    }
+    jobConfig.config.process[0].sample.samples = newSamples;
+    delete jobConfig.config.process[0].sample.prompts;
+  }
+  return jobConfig;
+};

src/app/jobs/new/options.ts ADDED Viewed

	@@ -0,0 +1,441 @@

+import { GroupedSelectOption, SelectOption } from '@/types';
+type Control = 'depth' | 'line' | 'pose' | 'inpaint';
+type DisableableSections = 'model.quantize' | 'train.timestep_type' | 'network.conv';
+type AdditionalSections =
+  | 'datasets.control_path'
+  | 'datasets.do_i2v'
+  | 'sample.ctrl_img'
+  | 'datasets.num_frames'
+  | 'model.multistage'
+  | 'model.low_vram';
+type ModelGroup = 'image' | 'instruction' | 'video';
+export interface ModelArch {
+  name: string;
+  label: string;
+  group: ModelGroup;
+  controls?: Control[];
+  isVideoModel?: boolean;
+  defaults?: { [key: string]: any };
+  disableSections?: DisableableSections[];
+  additionalSections?: AdditionalSections[];
+  accuracyRecoveryAdapters?: { [key: string]: string };
+}
+const defaultNameOrPath = '';
+export const modelArchs: ModelArch[] = [
+  {
+    name: 'flux',
+    label: 'FLUX.1',
+    group: 'image',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['black-forest-labs/FLUX.1-dev', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+    },
+    disableSections: ['network.conv'],
+  },
+  {
+    name: 'flux_kontext',
+    label: 'FLUX.1-Kontext-dev',
+    group: 'instruction',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['black-forest-labs/FLUX.1-Kontext-dev', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.timestep_type': ['weighted', 'sigmoid'],
+    },
+    disableSections: ['network.conv'],
+    additionalSections: ['datasets.control_path', 'sample.ctrl_img'],
+  },
+  {
+    name: 'flex1',
+    label: 'Flex.1',
+    group: 'image',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['ostris/Flex.1-alpha', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].train.bypass_guidance_embedding': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+    },
+    disableSections: ['network.conv'],
+  },
+  {
+    name: 'flex2',
+    label: 'Flex.2',
+    group: 'image',
+    controls: ['depth', 'line', 'pose', 'inpaint'],
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['ostris/Flex.2-preview', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].model.model_kwargs': [
+        {
+          invert_inpaint_mask_chance: 0.2,
+          inpaint_dropout: 0.5,
+          control_dropout: 0.5,
+          inpaint_random_chance: 0.2,
+          do_random_inpainting: true,
+          random_blur_mask: true,
+          random_dialate_mask: true,
+        },
+        {},
+      ],
+      'config.process[0].train.bypass_guidance_embedding': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+    },
+    disableSections: ['network.conv'],
+  },
+  {
+    name: 'chroma',
+    label: 'Chroma',
+    group: 'image',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['lodestones/Chroma1-Base', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+    },
+    disableSections: ['network.conv'],
+  },
+  {
+    name: 'wan21:1b',
+    label: 'Wan 2.1 (1.3B)',
+    group: 'video',
+    isVideoModel: true,
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['Wan-AI/Wan2.1-T2V-1.3B-Diffusers', defaultNameOrPath],
+      'config.process[0].model.quantize': [false, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].sample.num_frames': [41, 1],
+      'config.process[0].sample.fps': [16, 1],
+    },
+    disableSections: ['network.conv'],
+    additionalSections: ['datasets.num_frames', 'model.low_vram'],
+  },
+  {
+    name: 'wan21_i2v:14b480p',
+    label: 'Wan 2.1 I2V (14B-480P)',
+    group: 'video',
+    isVideoModel: true,
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['Wan-AI/Wan2.1-I2V-14B-480P-Diffusers', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].sample.num_frames': [41, 1],
+      'config.process[0].sample.fps': [16, 1],
+      'config.process[0].train.timestep_type': ['weighted', 'sigmoid'],
+    },
+    disableSections: ['network.conv'],
+    additionalSections: ['sample.ctrl_img', 'datasets.num_frames', 'model.low_vram'],
+  },
+  {
+    name: 'wan21_i2v:14b',
+    label: 'Wan 2.1 I2V (14B-720P)',
+    group: 'video',
+    isVideoModel: true,
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['Wan-AI/Wan2.1-I2V-14B-720P-Diffusers', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].sample.num_frames': [41, 1],
+      'config.process[0].sample.fps': [16, 1],
+      'config.process[0].train.timestep_type': ['weighted', 'sigmoid'],
+    },
+    disableSections: ['network.conv'],
+    additionalSections: ['sample.ctrl_img', 'datasets.num_frames', 'model.low_vram'],
+  },
+  {
+    name: 'wan21:14b',
+    label: 'Wan 2.1 (14B)',
+    group: 'video',
+    isVideoModel: true,
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['Wan-AI/Wan2.1-T2V-14B-Diffusers', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].sample.num_frames': [41, 1],
+      'config.process[0].sample.fps': [16, 1],
+    },
+    disableSections: ['network.conv'],
+    additionalSections: ['datasets.num_frames', 'model.low_vram'],
+  },
+  {
+    name: 'wan22_14b:t2v',
+    label: 'Wan 2.2 (14B)',
+    group: 'video',
+    isVideoModel: true,
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['ai-toolkit/Wan2.2-T2V-A14B-Diffusers-bf16', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].sample.num_frames': [41, 1],
+      'config.process[0].sample.fps': [16, 1],
+      'config.process[0].model.low_vram': [true, false],
+      'config.process[0].train.timestep_type': ['linear', 'sigmoid'],
+      'config.process[0].model.model_kwargs': [
+        {
+          train_high_noise: true,
+          train_low_noise: true,
+        },
+        {},
+      ],
+    },
+    disableSections: ['network.conv'],
+    additionalSections: ['datasets.num_frames', 'model.low_vram', 'model.multistage'],
+    accuracyRecoveryAdapters: {
+      // '3 bit with ARA': 'uint3|ostris/accuracy_recovery_adapters/wan22_14b_t2i_torchao_uint3.safetensors',
+      '4 bit with ARA': 'uint4|ostris/accuracy_recovery_adapters/wan22_14b_t2i_torchao_uint4.safetensors',
+    },
+  },
+  {
+    name: 'wan22_14b_i2v',
+    label: 'Wan 2.2 I2V (14B)',
+    group: 'video',
+    isVideoModel: true,
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['ai-toolkit/Wan2.2-I2V-A14B-Diffusers-bf16', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].sample.num_frames': [41, 1],
+      'config.process[0].sample.fps': [16, 1],
+      'config.process[0].model.low_vram': [true, false],
+      'config.process[0].train.timestep_type': ['linear', 'sigmoid'],
+      'config.process[0].model.model_kwargs': [
+        {
+          train_high_noise: true,
+          train_low_noise: true,
+        },
+        {},
+      ],
+    },
+    disableSections: ['network.conv'],
+    additionalSections: ['sample.ctrl_img', 'datasets.num_frames', 'model.low_vram', 'model.multistage'],
+    accuracyRecoveryAdapters: {
+      '4 bit with ARA': 'uint4|ostris/accuracy_recovery_adapters/wan22_14b_i2v_torchao_uint4.safetensors',
+    },
+  },
+  {
+    name: 'wan22_5b',
+    label: 'Wan 2.2 TI2V (5B)',
+    group: 'video',
+    isVideoModel: true,
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['Wan-AI/Wan2.2-TI2V-5B-Diffusers', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].model.low_vram': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].sample.num_frames': [121, 1],
+      'config.process[0].sample.fps': [24, 1],
+      'config.process[0].sample.width': [768, 1024],
+      'config.process[0].sample.height': [768, 1024],
+      'config.process[0].train.timestep_type': ['weighted', 'sigmoid'],
+    },
+    disableSections: ['network.conv'],
+    additionalSections: ['sample.ctrl_img', 'datasets.num_frames', 'model.low_vram', 'datasets.do_i2v'],
+  },
+  {
+    name: 'lumina2',
+    label: 'Lumina2',
+    group: 'image',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['Alpha-VLLM/Lumina-Image-2.0', defaultNameOrPath],
+      'config.process[0].model.quantize': [false, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+    },
+    disableSections: ['network.conv'],
+  },
+  {
+    name: 'qwen_image',
+    label: 'Qwen-Image',
+    group: 'image',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['Qwen/Qwen-Image', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].model.low_vram': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.timestep_type': ['weighted', 'sigmoid'],
+      'config.process[0].model.qtype': ['qfloat8', 'qfloat8'],
+    },
+    disableSections: ['network.conv'],
+    additionalSections: ['model.low_vram'],
+    accuracyRecoveryAdapters: {
+      '3 bit with ARA': 'uint3|ostris/accuracy_recovery_adapters/qwen_image_torchao_uint3.safetensors',
+    },
+  },
+  {
+    name: 'qwen_image_edit',
+    label: 'Qwen-Image-Edit',
+    group: 'instruction',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['Qwen/Qwen-Image-Edit', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].model.low_vram': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.timestep_type': ['weighted', 'sigmoid'],
+      'config.process[0].model.qtype': ['qfloat8', 'qfloat8'],
+    },
+    disableSections: ['network.conv'],
+    additionalSections: ['datasets.control_path', 'sample.ctrl_img', 'model.low_vram'],
+    accuracyRecoveryAdapters: {
+      '3 bit with ARA': 'uint3|ostris/accuracy_recovery_adapters/qwen_image_edit_torchao_uint3.safetensors',
+    },
+  },
+  {
+    name: 'hidream',
+    label: 'HiDream',
+    group: 'image',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['HiDream-ai/HiDream-I1-Full', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.lr': [0.0002, 0.0001],
+      'config.process[0].train.timestep_type': ['shift', 'sigmoid'],
+      'config.process[0].network.network_kwargs.ignore_if_contains': [['ff_i.experts', 'ff_i.gate'], []],
+    },
+    disableSections: ['network.conv'],
+    additionalSections: ['model.low_vram'],
+  },
+  {
+    name: 'hidream_e1',
+    label: 'HiDream E1',
+    group: 'instruction',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['HiDream-ai/HiDream-E1-1', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.lr': [0.0001, 0.0001],
+      'config.process[0].train.timestep_type': ['weighted', 'sigmoid'],
+      'config.process[0].network.network_kwargs.ignore_if_contains': [['ff_i.experts', 'ff_i.gate'], []],
+    },
+    disableSections: ['network.conv'],
+    additionalSections: ['datasets.control_path', 'sample.ctrl_img', 'model.low_vram'],
+  },
+  {
+    name: 'sdxl',
+    label: 'SDXL',
+    group: 'image',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['stabilityai/stable-diffusion-xl-base-1.0', defaultNameOrPath],
+      'config.process[0].model.quantize': [false, false],
+      'config.process[0].model.quantize_te': [false, false],
+      'config.process[0].sample.sampler': ['ddpm', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['ddpm', 'flowmatch'],
+      'config.process[0].sample.guidance_scale': [6, 4],
+    },
+    disableSections: ['model.quantize', 'train.timestep_type'],
+  },
+  {
+    name: 'sd15',
+    label: 'SD 1.5',
+    group: 'image',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['stable-diffusion-v1-5/stable-diffusion-v1-5', defaultNameOrPath],
+      'config.process[0].sample.sampler': ['ddpm', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['ddpm', 'flowmatch'],
+      'config.process[0].sample.width': [512, 1024],
+      'config.process[0].sample.height': [512, 1024],
+      'config.process[0].sample.guidance_scale': [6, 4],
+    },
+    disableSections: ['model.quantize', 'train.timestep_type'],
+  },
+  {
+    name: 'omnigen2',
+    label: 'OmniGen2',
+    group: 'image',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['OmniGen2/OmniGen2', defaultNameOrPath],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].model.quantize': [false, false],
+      'config.process[0].model.quantize_te': [true, false],
+    },
+    disableSections: ['network.conv'],
+    additionalSections: ['datasets.control_path', 'sample.ctrl_img'],
+  },
+].sort((a, b) => {
+  // Sort by label, case-insensitive
+  return a.label.localeCompare(b.label, undefined, { sensitivity: 'base' });
+}) as any;
+export const groupedModelOptions: GroupedSelectOption[] = modelArchs.reduce((acc, arch) => {
+  const group = acc.find(g => g.label === arch.group);
+  if (group) {
+    group.options.push({ value: arch.name, label: arch.label });
+  } else {
+    acc.push({
+      label: arch.group,
+      options: [{ value: arch.name, label: arch.label }],
+    });
+  }
+  return acc;
+}, [] as GroupedSelectOption[]);
+export const quantizationOptions: SelectOption[] = [
+  { value: '', label: '- NONE -' },
+  { value: 'qfloat8', label: 'float8 (default)' },
+  { value: 'uint8', label: '8 bit' },
+  { value: 'uint7', label: '7 bit' },
+  { value: 'uint6', label: '6 bit' },
+  { value: 'uint5', label: '5 bit' },
+  { value: 'uint4', label: '4 bit' },
+  { value: 'uint3', label: '3 bit' },
+  { value: 'uint2', label: '2 bit' },
+];
+export const defaultQtype = 'qfloat8';

src/app/jobs/new/page.tsx ADDED Viewed

	@@ -0,0 +1,306 @@

+'use client';
+import { useEffect, useState } from 'react';
+import { useSearchParams, useRouter } from 'next/navigation';
+import Link from 'next/link';
+import { defaultJobConfig, defaultDatasetConfig, migrateJobConfig } from './jobConfig';
+import { JobConfig } from '@/types';
+import { objectCopy } from '@/utils/basic';
+import { useNestedState } from '@/utils/hooks';
+import { SelectInput } from '@/components/formInputs';
+import useSettings from '@/hooks/useSettings';
+import useGPUInfo from '@/hooks/useGPUInfo';
+import useDatasetList from '@/hooks/useDatasetList';
+import path from 'path';
+import { TopBar, MainContent } from '@/components/layout';
+import { Button } from '@headlessui/react';
+import { FaChevronLeft } from 'react-icons/fa';
+import SimpleJob from './SimpleJob';
+import AdvancedJob from './AdvancedJob';
+import ErrorBoundary from '@/components/ErrorBoundary';
+import { getJob, upsertJob } from '@/utils/storage/jobStorage';
+import { usingBrowserDb } from '@/utils/env';
+import { getUserDatasetPath, updateUserDatasetPath } from '@/utils/storage/datasetStorage';
+import { apiClient } from '@/utils/api';
+import { useAuth } from '@/contexts/AuthContext';
+import HFLoginButton from '@/components/HFLoginButton';
+const isDev = process.env.NODE_ENV === 'development';
+export default function TrainingForm() {
+  const router = useRouter();
+  const searchParams = useSearchParams();
+  const runId = searchParams.get('id');
+  const { status: authStatus } = useAuth();
+  const isAuthenticated = authStatus === 'authenticated';
+  const [gpuIDs, setGpuIDs] = useState<string | null>(null);
+  const { settings, isSettingsLoaded } = useSettings();
+  const { gpuList, isGPUInfoLoaded } = useGPUInfo();
+  const { datasets, status: datasetFetchStatus } = useDatasetList();
+  const [datasetOptions, setDatasetOptions] = useState<{ value: string; label: string }[]>([]);
+  const [showAdvancedView, setShowAdvancedView] = useState(false);
+  const [jobConfig, setJobConfig] = useNestedState<JobConfig>(objectCopy(defaultJobConfig));
+  const [status, setStatus] = useState<'idle' | 'saving' | 'success' | 'error'>('idle');
+  // Track HF Jobs backend state
+  const [trainingBackend, setTrainingBackend] = useState<'local' | 'hf-jobs'>(
+    usingBrowserDb ? 'hf-jobs' : 'local',
+  );
+  const [hfJobSubmitted, setHfJobSubmitted] = useState(false);
+  useEffect(() => {
+    if (!isSettingsLoaded || !isAuthenticated) return;
+    if (datasetFetchStatus !== 'success') return;
+    let isMounted = true;
+    const buildDatasetOptions = async () => {
+      const options = await Promise.all(
+        datasets.map(async name => {
+          let datasetPath = settings.DATASETS_FOLDER ? path.join(settings.DATASETS_FOLDER, name) : '';
+          if (usingBrowserDb) {
+            const storedPath = getUserDatasetPath(name);
+            if (storedPath) {
+              datasetPath = storedPath;
+            } else {
+              try {
+                const response = await apiClient
+                  .post('/api/datasets/create', { name })
+                  .then(res => res.data);
+                if (response?.path) {
+                  datasetPath = response.path;
+                  updateUserDatasetPath(name, datasetPath);
+                }
+              } catch (err) {
+                console.error('Error resolving dataset path:', err);
+              }
+            }
+          }
+          if (!datasetPath) {
+            datasetPath = name;
+          }
+          return { value: datasetPath, label: name };
+        }),
+      );
+      if (!isMounted) {
+        return;
+      }
+      setDatasetOptions(options);
+      const defaultDatasetPath = defaultDatasetConfig.folder_path;
+      for (let i = 0; i < jobConfig.config.process[0].datasets.length; i++) {
+        const dataset = jobConfig.config.process[0].datasets[i];
+        if (dataset.folder_path === defaultDatasetPath) {
+          if (options.length > 0) {
+            setJobConfig(options[0].value, `config.process[0].datasets[${i}].folder_path`);
+          }
+        }
+      }
+    };
+    buildDatasetOptions();
+    return () => {
+      isMounted = false;
+    };
+  }, [datasets, settings, isSettingsLoaded, datasetFetchStatus]);
+  useEffect(() => {
+    if (runId) {
+      getJob(runId)
+        .then(data => {
+          if (!data) {
+            throw new Error('Job not found');
+          }
+          setGpuIDs(data.gpu_ids);
+          const parsedJobConfig = migrateJobConfig(JSON.parse(data.job_config));
+          setJobConfig(parsedJobConfig);
+          if (parsedJobConfig.is_hf_job) {
+            setTrainingBackend('hf-jobs');
+            setHfJobSubmitted(true);
+          }
+        })
+        .catch(error => console.error('Error fetching training:', error));
+    }
+  }, [runId]);
+  useEffect(() => {
+    if (isGPUInfoLoaded) {
+      if (gpuIDs === null && gpuList.length > 0) {
+        setGpuIDs(`${gpuList[0].index}`);
+      }
+    }
+  }, [gpuList, isGPUInfoLoaded]);
+  useEffect(() => {
+    if (isSettingsLoaded) {
+      setJobConfig(settings.TRAINING_FOLDER, 'config.process[0].training_folder');
+    }
+  }, [settings, isSettingsLoaded]);
+  const saveJob = async () => {
+    if (!isAuthenticated) return;
+    if (status === 'saving') return;
+    setStatus('saving');
+    try {
+      const savedJob = await upsertJob({
+        id: runId || undefined,
+        name: jobConfig.config.name,
+        gpu_ids: gpuIDs,
+        job_config: {
+          ...jobConfig,
+          is_hf_job: trainingBackend === 'hf-jobs',
+          hf_job_submitted: hfJobSubmitted,
+          training_backend: trainingBackend,
+        },
+        status: trainingBackend === 'hf-jobs' ? (hfJobSubmitted ? 'submitted' : 'stopped') : undefined,
+      });
+      setStatus('success');
+      router.push(`/jobs/${savedJob.id}`);
+    } catch (error: any) {
+      console.log('Error saving training:', error);
+      if (error?.code === 'P2002') {
+        alert('Training name already exists. Please choose a different name.');
+      } else {
+        alert('Failed to save job. Please try again.');
+      }
+    } finally {
+      setTimeout(() => {
+        setStatus('idle');
+      }, 2000);
+    }
+  };
+  const handleSubmit = async (e: React.FormEvent) => {
+    e.preventDefault();
+    saveJob();
+  };
+  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">{runId ? 'Edit Training Job' : 'New Training Job'}</h1>
+        </div>
+        <div className="flex-1"></div>
+        {showAdvancedView && isAuthenticated && (
+          <>
+            <div>
+              <SelectInput
+                value={`${gpuIDs}`}
+                onChange={value => setGpuIDs(value)}
+                options={gpuList.map((gpu: any) => ({ value: `${gpu.index}`, label: `GPU #${gpu.index}` }))}
+              />
+            </div>
+            <div className="mx-4 bg-gray-200 dark:bg-gray-800 w-1 h-6"></div>
+          </>
+        )}
+        <div className="pr-2">
+          <Button
+            className="text-gray-200 bg-gray-800 px-3 py-1 rounded-md"
+            onClick={() => setShowAdvancedView(!showAdvancedView)}
+          >
+            {showAdvancedView ? 'Show Simple' : 'Show Advanced'}
+          </Button>
+        </div>
+        <div>
+          <Button
+            className="text-gray-200 bg-green-800 hover:bg-green-700 px-3 py-1 rounded-md"
+            onClick={() => saveJob()}
+            disabled={!isAuthenticated || status === 'saving'}
+          >
+            {status === 'saving'
+              ? 'Saving...'
+              : runId
+              ? 'Update Job'
+              : 'Create Job'}
+          </Button>
+        </div>
+      </TopBar>
+      {!isAuthenticated ? (
+        <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 access token before creating or editing jobs.</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>
+      ) : showAdvancedView ? (
+        <div className="pt-[48px] absolute top-0 left-0 w-full h-full overflow-auto">
+          <AdvancedJob
+            jobConfig={jobConfig}
+            setJobConfig={setJobConfig}
+            status={status}
+            handleSubmit={handleSubmit}
+            runId={runId}
+            gpuIDs={gpuIDs}
+            setGpuIDs={setGpuIDs}
+            gpuList={gpuList}
+            datasetOptions={datasetOptions}
+            settings={settings}
+          />
+        </div>
+      ) : (
+        <MainContent>
+          <ErrorBoundary
+            fallback={
+              <div className="flex items-center justify-center h-64 text-lg text-red-600 font-medium bg-red-100 dark:bg-red-900/20 dark:text-red-400 border border-red-300 dark:border-red-700 rounded-lg">
+                Advanced job detected. Please switch to advanced view to continue.
+              </div>
+            }
+          >
+            <SimpleJob
+              jobConfig={jobConfig}
+              setJobConfig={setJobConfig}
+              status={status}
+              handleSubmit={handleSubmit}
+              runId={runId}
+              gpuIDs={gpuIDs}
+              setGpuIDs={setGpuIDs}
+              gpuList={gpuList}
+              datasetOptions={datasetOptions}
+              trainingBackend={trainingBackend}
+              setTrainingBackend={usingBrowserDb ? undefined : setTrainingBackend}
+              hfJobSubmitted={hfJobSubmitted}
+              onHFJobComplete={(jobId: string, localJobId?: string) => {
+                setHfJobSubmitted(true);
+                // Redirect to the job detail page
+                if (localJobId) {
+                  router.push(`/jobs/${localJobId}`);
+                }
+              }}
+              forceHFBackend={usingBrowserDb}
+            />
+          </ErrorBoundary>
+          <div className="pt-20"></div>
+        </MainContent>
+      )}
+    </>
+  );
+}
+  useEffect(() => {
+    if (!isAuthenticated) {
+      setDatasetOptions([]);
+    }
+  }, [isAuthenticated]);

src/app/jobs/page.tsx ADDED Viewed

	@@ -0,0 +1,49 @@

+'use client';
+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 } = useAuth();
+  const isAuthenticated = authStatus === 'authenticated';
+  return (
+    <>
+      <TopBar>
+        <div>
+          <h1 className="text-lg">Training Jobs</h1>
+        </div>
+        <div className="flex-1"></div>
+        <div>
+          {isAuthenticated ? (
+            <Link href="/jobs/new" className="text-gray-200 bg-slate-600 px-3 py-1 rounded-md">
+              New Training Job
+            </Link>
+          ) : (
+            <span className="text-gray-600 bg-gray-900 px-3 py-1 rounded-md border border-gray-800">
+              Sign in to create jobs
+            </span>
+          )}
+        </div>
+      </TopBar>
+      <MainContent>
+        {isAuthenticated ? (
+          <JobsTable />
+        ) : (
+          <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 a personal access token to view and manage training jobs.</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 tokens in Settings
+              </Link>
+            </div>
+          </div>
+        )}
+      </MainContent>
+    </>
+  );
+}

src/app/layout.tsx ADDED Viewed

	@@ -0,0 +1,50 @@

+import type { Metadata } from 'next';
+import { Inter } from 'next/font/google';
+import './globals.css';
+import Sidebar from '@/components/Sidebar';
+import { ThemeProvider } from '@/components/ThemeProvider';
+import ConfirmModal from '@/components/ConfirmModal';
+import SampleImageModal from '@/components/SampleImageModal';
+import { Suspense } from 'react';
+import AuthWrapper from '@/components/AuthWrapper';
+import DocModal from '@/components/DocModal';
+import { AuthProvider } from '@/contexts/AuthContext';
+export const dynamic = 'force-dynamic';
+const inter = Inter({ subsets: ['latin'] });
+export const metadata: Metadata = {
+  title: 'Ostris - AI Toolkit',
+  description: 'A toolkit for building AI things.',
+};
+export default function RootLayout({ children }: { children: React.ReactNode }) {
+  // Check if the AI_TOOLKIT_AUTH environment variable is set
+  const authRequired = process.env.AI_TOOLKIT_AUTH ? true : false;
+  return (
+    <html lang="en" className="dark">
+      <head>
+        <meta name="apple-mobile-web-app-title" content="AI-Toolkit" />
+      </head>
+      <body className={inter.className} suppressHydrationWarning={true}>
+        <ThemeProvider>
+          <AuthProvider>
+            <AuthWrapper authRequired={authRequired}>
+              <div className="flex h-screen bg-gray-950">
+                <Sidebar />
+                <main className="flex-1 overflow-auto bg-gray-950 text-gray-100 relative">
+                  <Suspense>{children}</Suspense>
+                </main>
+              </div>
+            </AuthWrapper>
+          </AuthProvider>
+        </ThemeProvider>
+        <ConfirmModal />
+        <DocModal />
+        <SampleImageModal />
+      </body>
+    </html>
+  );
+}

src/app/manifest.json ADDED Viewed

	@@ -0,0 +1,21 @@

+{
+  "name": "AI Toolkit",
+  "short_name": "AIToolkit",
+  "icons": [
+    {
+      "src": "/web-app-manifest-192x192.png",
+      "sizes": "192x192",
+      "type": "image/png",
+      "purpose": "maskable"
+    },
+    {
+      "src": "/web-app-manifest-512x512.png",
+      "sizes": "512x512",
+      "type": "image/png",
+      "purpose": "maskable"
+    }
+  ],
+  "theme_color": "#000000",
+  "background_color": "#000000",
+  "display": "standalone"
+}

src/app/page.tsx ADDED Viewed

	@@ -0,0 +1,5 @@

+import { redirect } from 'next/navigation';
+export default function Home() {
+  redirect('/dashboard');
+}

src/app/settings/page.tsx ADDED Viewed

	@@ -0,0 +1,264 @@

+'use client';
+import { useEffect, useState } from 'react';
+import useSettings from '@/hooks/useSettings';
+import { TopBar, MainContent } from '@/components/layout';
+import { persistSettings } from '@/utils/storage/settingsStorage';
+import { useAuth } from '@/contexts/AuthContext';
+import HFLoginButton from '@/components/HFLoginButton';
+import { useMemo } from 'react';
+import Link from 'next/link';
+export default function Settings() {
+  const { settings, setSettings } = useSettings();
+  const { status: authStatus, namespace, oauthAvailable, loginWithOAuth, logout, setManualToken, error: authError, token: authToken } = useAuth();
+  const [status, setStatus] = useState<'idle' | 'saving' | 'success' | 'error'>('idle');
+  const [manualToken, setManualTokenInput] = useState(settings.HF_TOKEN || '');
+  const isAuthenticated = authStatus === 'authenticated';
+  useEffect(() => {
+    setManualTokenInput(settings.HF_TOKEN || '');
+  }, [settings.HF_TOKEN]);
+  const handleSubmit = async (e: React.FormEvent) => {
+    e.preventDefault();
+    setStatus('saving');
+    persistSettings(settings)
+      .then(() => {
+        setStatus('success');
+      })
+      .catch(error => {
+        console.error('Error saving settings:', error);
+        setStatus('error');
+      })
+      .finally(() => {
+        setTimeout(() => setStatus('idle'), 2000);
+      });
+  };
+  const handleChange = (e: React.ChangeEvent<HTMLInputElement>) => {
+    const { name, value } = e.target;
+    setSettings(prev => ({ ...prev, [name]: value }));
+  };
+  const handleManualSubmit = async (e: React.FormEvent) => {
+    e.preventDefault();
+    await setManualToken(manualToken);
+  };
+  const authDescription = useMemo(() => {
+    if (authStatus === 'checking') {
+      return 'Checking your Hugging Face session…';
+    }
+    if (isAuthenticated) {
+      return `Connected as ${namespace}`;
+    }
+    return 'Sign in to use Hugging Face Jobs or submit your own access token.';
+  }, [authStatus, isAuthenticated, namespace]);
+  return (
+    <>
+      <TopBar>
+        <div>
+          <h1 className="text-lg">Settings</h1>
+        </div>
+        <div className="flex-1"></div>
+        <div className="flex items-center gap-3 pr-2 text-sm text-gray-400">
+          {isAuthenticated ? (
+            <span>Welcome, {namespace || 'user'}</span>
+          ) : (
+            <span>Authenticate to unlock training features</span>
+          )}
+        </div>
+      </TopBar>
+      <MainContent>
+        <div className="grid gap-4 md:grid-cols-2 mb-6">
+          <div className="border border-gray-800 rounded-xl p-5 bg-gray-900">
+            <div className="flex items-center justify-between mb-4">
+              <div>
+                <h2 className="text-md font-semibold text-gray-100">Sign in with Hugging Face</h2>
+                <p className="text-sm text-gray-400 mt-1">{authDescription}</p>
+              </div>
+              {isAuthenticated && (
+                <span className="text-xs px-2 py-1 rounded-full bg-emerald-900 text-emerald-300">Authenticated</span>
+              )}
+            </div>
+            <div className="flex items-center gap-3">
+              {isAuthenticated ? (
+                <button
+                  type="button"
+                  onClick={logout}
+                  className="px-4 py-2 rounded-md border border-gray-700 text-sm bg-gray-800 hover:bg-gray-700 transition-colors"
+                >
+                  Sign out
+                </button>
+              ) : (
+                <>
+                  <HFLoginButton size="md" className="bg-transparent border-none p-0" />
+                  {!oauthAvailable && (
+                    <span className="text-xs text-yellow-500">
+                      OAuth is unavailable. Set HF_OAUTH_CLIENT_ID/SECRET on the server.
+                    </span>
+                  )}
+                </>
+              )}
+            </div>
+            {!isAuthenticated && authError && (
+              <p className="mt-3 text-xs text-red-400">{authError}</p>
+            )}
+          </div>
+          <form onSubmit={handleManualSubmit} className="border border-gray-800 rounded-xl p-5 bg-gray-900">
+            <h2 className="text-md font-semibold text-gray-100">Manual Token</h2>
+            <p className="text-sm text-gray-400 mt-1">
+              Paste an access token created at{' '}
+              <a href="https://huggingface.co/settings/tokens" target="_blank" rel="noreferrer" className="text-blue-400 hover:text-blue-300">
+                huggingface.co/settings/tokens
+              </a>
+              .
+            </p>
+            <div className="mt-4">
+              <input
+                type="password"
+                value={manualToken}
+                onChange={event => setManualTokenInput(event.target.value)}
+                className="w-full px-4 py-2 bg-gray-800 border border-gray-700 rounded-lg focus:ring-2 focus:ring-gray-600 focus:border-transparent"
+                placeholder="Enter Hugging Face token"
+              />
+            </div>
+            <div className="mt-4 flex items-center gap-3">
+              <button
+                type="submit"
+                className="px-4 py-2 rounded-md bg-blue-600 hover:bg-blue-500 text-sm text-white transition-colors disabled:opacity-50 disabled:cursor-not-allowed"
+                disabled={authStatus === 'checking' || manualToken.trim() === ''}
+              >
+                Validate Token
+              </button>
+              {isAuthenticated && authToken === manualToken && (
+                <span className="text-xs text-emerald-400">Active token</span>
+              )}
+            </div>
+            {authError && (
+              <p className="mt-3 text-xs text-red-400">{authError}</p>
+            )}
+          </form>
+        </div>
+        <form onSubmit={handleSubmit} className="space-y-6">
+          <div className="grid grid-cols-1 gap-6 sm:grid-cols-2">
+            <div>
+              <div className="space-y-4">
+                <div>
+                  <label htmlFor="TRAINING_FOLDER" className="block text-sm font-medium mb-2">
+                    Training Folder Path
+                    <div className="text-gray-500 text-sm ml-1">
+                      We will store your training information here. Must be an absolute path. If blank, it will default
+                      to the output folder in the project root.
+                    </div>
+                  </label>
+                  <input
+                    type="text"
+                    id="TRAINING_FOLDER"
+                    name="TRAINING_FOLDER"
+                    value={settings.TRAINING_FOLDER}
+                    onChange={handleChange}
+                    className="w-full px-4 py-2 bg-gray-800 border border-gray-700 rounded-lg focus:ring-2 focus:ring-gray-600 focus:border-transparent"
+                    placeholder="Enter training folder path"
+                  />
+                </div>
+                <div>
+                  <label htmlFor="DATASETS_FOLDER" className="block text-sm font-medium mb-2">
+                    Dataset Folder Path
+                    <div className="text-gray-500 text-sm ml-1">
+                      Where we store and find your datasets.{' '}
+                      <span className="text-orange-800">
+                        Warning: This software may modify datasets so it is recommended you keep a backup somewhere else
+                        or have a dedicated folder for this software.
+                      </span>
+                    </div>
+                  </label>
+                  <input
+                    type="text"
+                    id="DATASETS_FOLDER"
+                    name="DATASETS_FOLDER"
+                    value={settings.DATASETS_FOLDER}
+                    onChange={handleChange}
+                    className="w-full px-4 py-2 bg-gray-800 border border-gray-700 rounded-lg focus:ring-2 focus:ring-gray-600 focus:border-transparent"
+                    placeholder="Enter datasets folder path"
+                  />
+                </div>
+              </div>
+            </div>
+            <div>
+              <div className="space-y-4">
+                <h3 className="text-lg font-medium mb-4">Hugging Face Jobs (Cloud Training)</h3>
+                <div>
+                  <label htmlFor="HF_JOBS_NAMESPACE" className="block text-sm font-medium mb-2">
+                    HF Jobs Namespace (optional)
+                    <div className="text-gray-500 text-sm ml-1">
+                      Leave blank to default to the account associated with your Hugging Face token.
+                    </div>
+                  </label>
+                  <input
+                    type="text"
+                    id="HF_JOBS_NAMESPACE"
+                    name="HF_JOBS_NAMESPACE"
+                    value={settings.HF_JOBS_NAMESPACE}
+                    onChange={handleChange}
+                    className="w-full px-4 py-2 bg-gray-800 border border-gray-700 rounded-lg focus:ring-2 focus:ring-gray-600 focus:border-transparent"
+                    placeholder="e.g. your-username or your-org"
+                  />
+                </div>
+                <div>
+                  <label htmlFor="HF_JOBS_DEFAULT_HARDWARE" className="block text-sm font-medium mb-2">
+                    Default Hardware
+                    <div className="text-gray-500 text-sm ml-1">
+                      Default hardware configuration for cloud training jobs.
+                    </div>
+                  </label>
+                  <select
+                    id="HF_JOBS_DEFAULT_HARDWARE"
+                    name="HF_JOBS_DEFAULT_HARDWARE"
+                    value={settings.HF_JOBS_DEFAULT_HARDWARE}
+                    onChange={(e) => setSettings(prev => ({ ...prev, HF_JOBS_DEFAULT_HARDWARE: e.target.value }))}
+                    className="w-full px-4 py-2 bg-gray-800 border border-gray-700 rounded-lg focus:ring-2 focus:ring-gray-600 focus:border-transparent"
+                  >
+                    <option value="cpu-basic">CPU Basic</option>
+                    <option value="cpu-upgrade">CPU Upgrade</option>
+                    <option value="t4-small">T4 Small</option>
+                    <option value="t4-medium">T4 Medium</option>
+                    <option value="l4x1">L4x1</option>
+                    <option value="l4x4">L4x4</option>
+                    <option value="a10g-small">A10G Small</option>
+                    <option value="a10g-large">A10G Large</option>
+                    <option value="a10g-largex2">A10G Large x2</option>
+                    <option value="a10g-largex4">A10G Large x4</option>
+                    <option value="a100-large">A100 Large</option>
+                    <option value="v5e-1x1">TPU v5e-1x1</option>
+                    <option value="v5e-2x2">TPU v5e-2x2</option>
+                    <option value="v5e-2x4">TPU v5e-2x4</option>
+                  </select>
+                </div>
+              </div>
+            </div>
+          </div>
+          <button
+            type="submit"
+            disabled={status === 'saving'}
+            className="w-full px-4 py-2 bg-gray-700 hover:bg-gray-600 rounded-lg transition-colors disabled:opacity-50 disabled:cursor-not-allowed"
+          >
+            {status === 'saving' ? 'Saving...' : 'Save Settings'}
+          </button>
+          {status === 'success' && <p className="text-green-500 text-center">Settings saved successfully!</p>}
+          {status === 'error' && <p className="text-red-500 text-center">Error saving settings. Please try again.</p>}
+        </form>
+      </MainContent>
+    </>
+  );
+}

src/components/AddImagesModal.tsx ADDED Viewed

	@@ -0,0 +1,152 @@

+'use client';
+import { createGlobalState } from 'react-global-hooks';
+import { Dialog, DialogBackdrop, DialogPanel, DialogTitle } from '@headlessui/react';
+import { FaUpload } from 'react-icons/fa';
+import { useCallback, useState } from 'react';
+import { useDropzone } from 'react-dropzone';
+import { apiClient } from '@/utils/api';
+export interface AddImagesModalState {
+  datasetName: string;
+  onComplete?: () => void;
+}
+export const addImagesModalState = createGlobalState<AddImagesModalState | null>(null);
+export const openImagesModal = (datasetName: string, onComplete: () => void) => {
+  addImagesModalState.set({ datasetName, onComplete });
+};
+export default function AddImagesModal() {
+  const [addImagesModalInfo, setAddImagesModalInfo] = addImagesModalState.use();
+  const [uploadProgress, setUploadProgress] = useState<number>(0);
+  const [isUploading, setIsUploading] = useState<boolean>(false);
+  const open = addImagesModalInfo !== null;
+  const onCancel = () => {
+    if (!isUploading) {
+      setAddImagesModalInfo(null);
+    }
+  };
+  const onDone = () => {
+    if (addImagesModalInfo?.onComplete && !isUploading) {
+      addImagesModalInfo.onComplete();
+      setAddImagesModalInfo(null);
+    }
+  };
+  const onDrop = useCallback(
+    async (acceptedFiles: File[]) => {
+      if (acceptedFiles.length === 0) return;
+      setIsUploading(true);
+      setUploadProgress(0);
+      const formData = new FormData();
+      acceptedFiles.forEach(file => {
+        formData.append('files', file);
+      });
+      formData.append('datasetName', addImagesModalInfo?.datasetName || '');
+      try {
+        await apiClient.post(`/api/datasets/upload`, formData, {
+          headers: {
+            'Content-Type': 'multipart/form-data',
+          },
+          onUploadProgress: progressEvent => {
+            const percentCompleted = Math.round((progressEvent.loaded * 100) / (progressEvent.total || 100));
+            setUploadProgress(percentCompleted);
+          },
+          timeout: 0, // Disable timeout
+        });
+        onDone();
+      } catch (error) {
+        console.error('Upload failed:', error);
+      } finally {
+        setIsUploading(false);
+        setUploadProgress(0);
+      }
+    },
+    [addImagesModalInfo],
+  );
+  const { getRootProps, getInputProps, isDragActive } = useDropzone({
+    onDrop,
+    accept: {
+      'image/*': ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp'],
+      'video/*': ['.mp4', '.avi', '.mov', '.mkv', '.wmv', '.m4v', '.flv'],
+      'text/*': ['.txt'],
+    },
+    multiple: true,
+  });
+  return (
+    <Dialog open={open} onClose={onCancel} className="relative z-10">
+      <DialogBackdrop
+        transition
+        className="fixed inset-0 bg-gray-900/75 transition-opacity data-closed:opacity-0 data-enter:duration-300 data-enter:ease-out data-leave:duration-200 data-leave:ease-in"
+      />
+      <div className="fixed inset-0 z-10 w-screen overflow-y-auto">
+        <div className="flex min-h-full items-end justify-center p-4 text-center sm:items-center sm:p-0">
+          <DialogPanel
+            transition
+            className="relative transform overflow-hidden rounded-lg bg-gray-800 text-left shadow-xl transition-all data-closed:translate-y-4 data-closed:opacity-0 data-enter:duration-300 data-enter:ease-out data-leave:duration-200 data-leave:ease-in sm:my-8 sm:w-full sm:max-w-lg data-closed:sm:translate-y-0 data-closed:sm:scale-95"
+          >
+            <div className="bg-gray-800 px-4 pt-5 pb-4 sm:p-6 sm:pb-4">
+              <div className="text-center">
+                <DialogTitle as="h3" className="text-base font-semibold text-gray-200 mb-4">
+                  Add Images to: {addImagesModalInfo?.datasetName}
+                </DialogTitle>
+                <div className="w-full">
+                  <div
+                    {...getRootProps()}
+                    className={`h-40 w-full flex flex-col items-center justify-center border-2 border-dashed rounded-lg cursor-pointer transition-colors duration-200
+                      ${isDragActive ? 'border-blue-500 bg-blue-50/10' : 'border-gray-600'}`}
+                  >
+                    <input {...getInputProps()} />
+                    <FaUpload className="size-8 mb-3 text-gray-400" />
+                    <p className="text-sm text-gray-200 text-center">
+                      {isDragActive ? 'Drop the files here...' : 'Drag & drop files here, or click to select files'}
+                    </p>
+                  </div>
+                  {isUploading && (
+                    <div className="mt-4">
+                      <div className="w-full bg-gray-700 rounded-full h-2.5">
+                        <div className="bg-blue-600 h-2.5 rounded-full" style={{ width: `${uploadProgress}%` }}></div>
+                      </div>
+                      <p className="text-sm text-gray-300 mt-2 text-center">Uploading... {uploadProgress}%</p>
+                    </div>
+                  )}
+                </div>
+              </div>
+            </div>
+            <div className="bg-gray-700 px-4 py-3 sm:flex sm:flex-row-reverse sm:px-6">
+              <button
+                type="button"
+                onClick={onDone}
+                disabled={isUploading}
+                className={`inline-flex w-full justify-center rounded-md bg-slate-600 px-3 py-2 text-sm font-semibold text-white shadow-xs sm:ml-3 sm:w-auto
+                  ${isUploading ? 'opacity-50 cursor-not-allowed' : ''}`}
+              >
+                Done
+              </button>
+              <button
+                type="button"
+                data-autofocus
+                onClick={onCancel}
+                disabled={isUploading}
+                className={`mt-3 inline-flex w-full justify-center rounded-md bg-gray-800 px-3 py-2 text-sm font-semibold text-gray-200 hover:bg-gray-800 sm:mt-0 sm:w-auto ring-0
+                  ${isUploading ? 'opacity-50 cursor-not-allowed' : ''}`}
+              >
+                Cancel
+              </button>
+            </div>
+          </DialogPanel>
+        </div>
+      </div>
+    </Dialog>
+  );
+}

src/components/AddSingleImageModal.tsx ADDED Viewed

	@@ -0,0 +1,141 @@

+'use client';
+import { createGlobalState } from 'react-global-hooks';
+import { Dialog, DialogBackdrop, DialogPanel, DialogTitle } from '@headlessui/react';
+import { FaUpload } from 'react-icons/fa';
+import { useCallback, useState } from 'react';
+import { useDropzone } from 'react-dropzone';
+import { apiClient } from '@/utils/api';
+export interface AddSingleImageModalState {
+  onComplete?: (imagePath: string|null) => void;
+}
+export const addSingleImageModalState = createGlobalState<AddSingleImageModalState | null>(null);
+export const openAddImageModal = (onComplete: (imagePath: string|null) => void) => {
+  addSingleImageModalState.set({onComplete });
+};
+export default function AddSingleImageModal() {
+  const [addSingleImageModalInfo, setAddSingleImageModalInfo] = addSingleImageModalState.use();
+  const [uploadProgress, setUploadProgress] = useState<number>(0);
+  const [isUploading, setIsUploading] = useState<boolean>(false);
+  const open = addSingleImageModalInfo !== null;
+  const onCancel = () => {
+    if (!isUploading) {
+      setAddSingleImageModalInfo(null);
+    }
+  };
+  const onDone = (imagePath: string|null) => {
+    if (addSingleImageModalInfo?.onComplete && !isUploading) {
+      addSingleImageModalInfo.onComplete(imagePath);
+      setAddSingleImageModalInfo(null);
+    }
+  };
+  const onDrop = useCallback(
+    async (acceptedFiles: File[]) => {
+      if (acceptedFiles.length === 0) return;
+      setIsUploading(true);
+      setUploadProgress(0);
+      const formData = new FormData();
+      acceptedFiles.forEach(file => {
+        formData.append('files', file);
+      });
+      try {
+        const resp = await apiClient.post(`/api/img/upload`, formData, {
+          headers: {
+            'Content-Type': 'multipart/form-data',
+          },
+          onUploadProgress: progressEvent => {
+            const percentCompleted = Math.round((progressEvent.loaded * 100) / (progressEvent.total || 100));
+            setUploadProgress(percentCompleted);
+          },
+          timeout: 0, // Disable timeout
+        });
+        console.log('Upload successful:', resp.data);
+        onDone(resp.data.files[0] || null);
+      } catch (error) {
+        console.error('Upload failed:', error);
+      } finally {
+        setIsUploading(false);
+        setUploadProgress(0);
+      }
+    },
+    [addSingleImageModalInfo],
+  );
+  const { getRootProps, getInputProps, isDragActive } = useDropzone({
+    onDrop,
+    accept: {
+      'image/*': ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp'],
+    },
+    multiple: false,
+  });
+  return (
+    <Dialog open={open} onClose={onCancel} className="relative z-10">
+      <DialogBackdrop
+        transition
+        className="fixed inset-0 bg-gray-900/75 transition-opacity data-closed:opacity-0 data-enter:duration-300 data-enter:ease-out data-leave:duration-200 data-leave:ease-in"
+      />
+      <div className="fixed inset-0 z-10 w-screen overflow-y-auto">
+        <div className="flex min-h-full items-end justify-center p-4 text-center sm:items-center sm:p-0">
+          <DialogPanel
+            transition
+            className="relative transform overflow-hidden rounded-lg bg-gray-800 text-left shadow-xl transition-all data-closed:translate-y-4 data-closed:opacity-0 data-enter:duration-300 data-enter:ease-out data-leave:duration-200 data-leave:ease-in sm:my-8 sm:w-full sm:max-w-lg data-closed:sm:translate-y-0 data-closed:sm:scale-95"
+          >
+            <div className="bg-gray-800 px-4 pt-5 pb-4 sm:p-6 sm:pb-4">
+              <div className="text-center">
+                <DialogTitle as="h3" className="text-base font-semibold text-gray-200 mb-4">
+                  Add Control Image
+                </DialogTitle>
+                <div className="w-full">
+                  <div
+                    {...getRootProps()}
+                    className={`h-40 w-full flex flex-col items-center justify-center border-2 border-dashed rounded-lg cursor-pointer transition-colors duration-200
+                      ${isDragActive ? 'border-blue-500 bg-blue-50/10' : 'border-gray-600'}`}
+                  >
+                    <input {...getInputProps()} />
+                    <FaUpload className="size-8 mb-3 text-gray-400" />
+                    <p className="text-sm text-gray-200 text-center">
+                      {isDragActive ? 'Drop the image here...' : 'Drag & drop an image here, or click to select one'}
+                    </p>
+                  </div>
+                  {isUploading && (
+                    <div className="mt-4">
+                      <div className="w-full bg-gray-700 rounded-full h-2.5">
+                        <div className="bg-blue-600 h-2.5 rounded-full" style={{ width: `${uploadProgress}%` }}></div>
+                      </div>
+                      <p className="text-sm text-gray-300 mt-2 text-center">Uploading... {uploadProgress}%</p>
+                    </div>
+                  )}
+                </div>
+              </div>
+            </div>
+            <div className="bg-gray-700 px-4 py-3 sm:flex sm:flex-row-reverse sm:px-6">
+              <button
+                type="button"
+                data-autofocus
+                onClick={onCancel}
+                disabled={isUploading}
+                className={`mt-3 inline-flex w-full justify-center rounded-md bg-gray-800 px-3 py-2 text-sm font-semibold text-gray-200 hover:bg-gray-800 sm:mt-0 sm:w-auto ring-0
+                  ${isUploading ? 'opacity-50 cursor-not-allowed' : ''}`}
+              >
+                Cancel
+              </button>
+            </div>
+          </DialogPanel>
+        </div>
+      </div>
+    </Dialog>
+  );
+}

src/components/AuthWrapper.tsx ADDED Viewed

	@@ -0,0 +1,166 @@

+'use client';
+import { useState, useEffect, useRef } from 'react';
+import { apiClient, isAuthorizedState } from '@/utils/api';
+import { createGlobalState } from 'react-global-hooks';
+interface AuthWrapperProps {
+  authRequired: boolean;
+  children: React.ReactNode | React.ReactNode[];
+}
+export default function AuthWrapper({ authRequired, children }: AuthWrapperProps) {
+  const [token, setToken] = useState('');
+  // start with true, and deauth if needed
+  const [isAuthorizedGlobal, setIsAuthorized] = isAuthorizedState.use();
+  const [isLoading, setIsLoading] = useState(false);
+  const [error, setError] = useState('');
+  const [isBrowser, setIsBrowser] = useState(false);
+  const inputRef = useRef<HTMLInputElement>(null);
+  const isAuthorized = authRequired ? isAuthorizedGlobal : true;
+  // Set isBrowser to true when component mounts
+  useEffect(() => {
+    setIsBrowser(true);
+    // Get token from localStorage only after component has mounted
+    const storedToken = localStorage.getItem('AI_TOOLKIT_AUTH') || '';
+    setToken(storedToken);
+    checkAuth();
+  }, []);
+  // auto focus on input when not authorized
+  useEffect(() => {
+    if (isAuthorized) {
+      return;
+    }
+    setTimeout(() => {
+      if (inputRef.current) {
+        inputRef.current.focus();
+      }
+    }, 100);
+  }, [isAuthorized]);
+  const checkAuth = async () => {
+    // always get current stored token here to avoid state race conditions
+    const currentToken = localStorage.getItem('AI_TOOLKIT_AUTH') || '';
+    if (!authRequired || isLoading || currentToken === '') {
+      return;
+    }
+    setIsLoading(true);
+    setError('');
+    try {
+      const response = await apiClient.get('/api/auth');
+      if (response.data.isAuthenticated) {
+        setIsAuthorized(true);
+      } else {
+        setIsAuthorized(false);
+        setError('Invalid token. Please try again.');
+      }
+    } catch (err) {
+      setIsAuthorized(false);
+      console.log(err);
+      setError('Invalid token. Please try again.');
+    }
+    setIsLoading(false);
+  };
+  const handleSubmit = async (e: React.FormEvent) => {
+    e.preventDefault();
+    setError('');
+    if (!token.trim()) {
+      setError('Please enter your token');
+      return;
+    }
+    if (isBrowser) {
+      localStorage.setItem('AI_TOOLKIT_AUTH', token);
+      checkAuth();
+    }
+  };
+  if (isAuthorized) {
+    return <>{children}</>;
+  }
+  return (
+    <div className="flex min-h-screen bg-gray-900 text-gray-100 absolute top-0 left-0 right-0 bottom-0 scroll-auto">
+      {/* Left side - decorative or brand area */}
+      <div className="hidden lg:flex lg:w-1/2 bg-gray-800 flex-col justify-center items-center p-12">
+        <div className="mb-4">
+          {/* Replace with your own logo */}
+          <div className="flex items-center justify-center">
+            <img src="/ostris_logo.png" alt="Ostris AI Toolkit" className="w-auto h-24 inline" />
+          </div>
+        </div>
+        <h1 className="text-4xl mb-6">AI Toolkit</h1>
+      </div>
+      {/* Right side - login form */}
+      <div className="w-full lg:w-1/2 flex flex-col justify-center items-center p-8 sm:p-12">
+        <div className="w-full max-w-md">
+          <div className="lg:hidden flex justify-center mb-4">
+            {/* Mobile logo */}
+            <div className="flex items-center justify-center">
+              <img src="/ostris_logo.png" alt="Ostris AI Toolkit" className="w-auto h-24 inline" />
+            </div>
+          </div>
+          <h2 className="text-3xl text-center mb-2 lg:hidden">AI Toolkit</h2>
+          <form onSubmit={handleSubmit} className="space-y-6">
+            <div>
+              <label htmlFor="token" className="block text-sm font-medium text-gray-400 mb-2">
+                Password
+              </label>
+              <input
+                id="token"
+                name="token"
+                type="password"
+                autoComplete="off"
+                required
+                value={token}
+                ref={inputRef}
+                onChange={e => setToken(e.target.value)}
+                className="w-full px-4 py-3 rounded-lg bg-gray-800 border border-gray-700 focus:border-blue-500 focus:ring-2 focus:ring-blue-500 focus:ring-opacity-50 text-gray-100 transition duration-200"
+                placeholder="Enter your password"
+              />
+              <div className='text-gray-500 text-xs mt-2'>
+                The password is set with the environment variable AI_TOOLKIT_AUTH, the default is the super secure secret word "password"
+              </div>
+            </div>
+            {error && (
+              <div className="p-3 bg-red-900/50 border border-red-800 rounded-lg text-red-200 text-sm">{error}</div>
+            )}
+            <button
+              type="submit"
+              disabled={isLoading}
+              className="w-full py-3 px-4 bg-blue-600 hover:bg-blue-700 rounded-lg text-white font-medium focus:outline-none focus:ring-2 focus:ring-blue-500 focus:ring-opacity-50 transition duration-200 flex items-center justify-center"
+            >
+              {isLoading ? (
+                <svg
+                  className="animate-spin h-5 w-5 text-white"
+                  xmlns="http://www.w3.org/2000/svg"
+                  fill="none"
+                  viewBox="0 0 24 24"
+                >
+                  <circle className="opacity-25" cx="12" cy="12" r="10" stroke="currentColor" strokeWidth="4"></circle>
+                  <path
+                    className="opacity-75"
+                    fill="currentColor"
+                    d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4zm2 5.291A7.962 7.962 0 014 12H0c0 3.042 1.135 5.824 3 7.938l3-2.647z"
+                  ></path>
+                </svg>
+              ) : (
+                'Check Password'
+              )}
+            </button>
+          </form>
+        </div>
+      </div>
+    </div>
+  );
+}

src/components/Card.tsx ADDED Viewed

	@@ -0,0 +1,15 @@

+interface CardProps {
+  title?: string;
+  children?: React.ReactNode;
+}
+const Card: React.FC<CardProps> = ({ title, children }) => {
+  return (
+    <section className="space-y-2 px-4 pb-4 pt-2 bg-gray-900 rounded-lg">
+      {title && <h2 className="text-lg mb-2 font-semibold uppercase text-gray-500">{title}</h2>}
+      {children ? children : null}
+    </section>
+  );
+};
+export default Card;