More ui more ui

.gitignore

@@ -161,6 +161,7 @@ cython_debug/
 
 /env.sh
 /models
+/datasets
 /custom/*
 !/custom/.gitkeep
 /.tmp

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

@@ -0,0 +1,42 @@
+// src/app/api/img/[imagePath]/route.ts
+import { NextRequest, NextResponse } from 'next/server';
+import fs from 'fs';
+import path from 'path';
+import { getDatasetsRoot } from '@/app/api/datasets/utils';
+
+export async function GET(request: NextRequest, { params }: { params: { imagePath: string } }) {
+  const { imagePath } = await params;
+  try {
+    // Decode the path
+    const filepath = decodeURIComponent(imagePath);
+
+    // caption name is the filepath without extension but with .txt
+    const captionPath = filepath.replace(/\.[^/.]+$/, '') + '.txt';
+
+    // Get allowed directories
+    const allowedDir = await getDatasetsRoot();
+
+    // Security check: Ensure path is in allowed directory
+    const isAllowed = filepath.startsWith(allowedDir) && !filepath.includes('..');
+
+    if (!isAllowed) {
+      console.warn(`Access denied: ${filepath} not in ${allowedDir}`);
+      return new NextResponse('Access denied', { status: 403 });
+    }
+
+    // Check if file exists
+    if (!fs.existsSync(captionPath)) {
+      // send back blank string if caption file does not exist
+      return new NextResponse('');
+    }
+
+    // Read caption file
+    const caption = fs.readFileSync(captionPath, 'utf-8');
+
+    // Return caption
+    return new NextResponse(caption);
+  } catch (error) {
+    console.error('Error getting caption:', error);
+    return new NextResponse('Error getting caption', { status: 500 });
+  }
+}

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

@@ -0,0 +1,22 @@
+import { NextResponse } from 'next/server';
+import fs from 'fs';
+import path from 'path';
+import { getDatasetsRoot } from '@/app/api/datasets/utils';
+
+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, create it
+    if (!fs.existsSync(datasetPath)) {
+      fs.mkdirSync(datasetPath);
+    }
+
+    return NextResponse.json({ success: true });
+  } catch (error) {
+    return NextResponse.json({ error: 'Failed to create dataset' }, { status: 500 });
+  }
+}

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

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

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

@@ -0,0 +1,67 @@
+import { NextResponse } from 'next/server';
+import fs from 'fs';
+import path from 'path';
+import { getDatasetsRoot } from '@/app/api/datasets/utils';
+
+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'];
+  let results: string[] = [];
+  
+  const items = fs.readdirSync(dir);
+  
+  for (const item of items) {
+    const itemPath = path.join(dir, item);
+    const stat = fs.statSync(itemPath);
+    
+    if (stat.isDirectory()) {
+      // If it's a directory, recursively search it
+      results = results.concat(findImagesRecursively(itemPath));
+    } else {
+      // If it's a file, check if it's an image
+      const ext = path.extname(itemPath).toLowerCase();
+      if (imageExtensions.includes(ext)) {
+        results.push(itemPath);
+      }
+    }
+  }
+  
+  return results;
+}

ui/src/app/api/datasets/utils.ts

@@ -0,0 +1,17 @@
+import { PrismaClient } from '@prisma/client';
+import { defaultDatasetsFolder } from '@/paths';
+
+const prisma = new PrismaClient();
+
+export const getDatasetsRoot = async () => {
+  let row = await prisma.settings.findFirst({
+    where: {
+      key: 'DATASETS_FOLDER',
+    },
+  });
+  let datasetsPath = defaultDatasetsFolder;
+  if (row?.value && row.value !== '') {
+    datasetsPath = row.value;
+  }
+  return datasetsPath;
+};

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

@@ -0,0 +1,66 @@
+// src/app/api/img/[imagePath]/route.ts
+import { NextRequest, NextResponse } from 'next/server';
+import fs from 'fs';
+import path from 'path';
+import { getDatasetsRoot } from '@/app/api/datasets/utils';
+
+export async function GET(request: NextRequest, { params }: { params: { imagePath: string } }) {
+  const { imagePath } = await params;
+  try {
+    // Decode the path
+    const filepath = decodeURIComponent(imagePath);
+    console.log('Serving image:', filepath);
+
+    // 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(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 } = {
+      '.jpg': 'image/jpeg',
+      '.jpeg': 'image/jpeg',
+      '.png': 'image/png',
+      '.gif': 'image/gif',
+      '.webp': 'image/webp',
+      '.svg': 'image/svg+xml',
+      '.bmp': 'image/bmp',
+    };
+
+    const contentType = contentTypeMap[ext] || 'application/octet-stream';
+
+    // Read file as buffer
+    const fileBuffer = fs.readFileSync(filepath);
+
+    // Return file with appropriate headers
+    return new NextResponse(fileBuffer, {
+      headers: {
+        'Content-Type': contentType,
+        'Content-Length': String(stat.size),
+        'Cache-Control': 'public, max-age=86400',
+      },
+    });
+  } catch (error) {
+    console.error('Error serving image:', error);
+    return new NextResponse('Internal Server Error', { status: 500 });
+  }
+}

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

@@ -1,6 +1,6 @@
 import { NextResponse } from 'next/server';
 import { PrismaClient } from '@prisma/client';
-import { defaultTrainFolder } from '@/paths';
+import { defaultTrainFolder, defaultDatasetsFolder } from '@/paths';
 
 const prisma = new PrismaClient();
 
@@ -15,6 +15,10 @@ export async function GET() {
     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 });
@@ -24,7 +28,7 @@ export async function GET() {
 export async function POST(request: Request) {
   try {
     const body = await request.json();
-    const { HF_TOKEN, TRAINING_FOLDER } = body;
+    const { HF_TOKEN, TRAINING_FOLDER, DATASETS_FOLDER } = body;
 
     // Upsert both settings
     await Promise.all([
@@ -38,6 +42,11 @@ export async function POST(request: Request) {
         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 },
+      }),
     ]);
 
     return NextResponse.json({ success: true });

ui/src/app/dashboard/page.tsx

@@ -5,7 +5,7 @@ import GpuMonitor from '@/components/GPUMonitor';
 export default function Dashboard() {
   return (
     <div className="space-y-6">
-      <h1 className="text-3xl font-bold">Dashboard</h1>
+      <h1 className="text-xl font-bold mb-8">Dashboard</h1>
       <GpuMonitor />
     </div>
   );

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

@@ -0,0 +1,73 @@
+'use client';
+
+import { useEffect, useState, use } from 'react';
+import Card from '@/components/Card';
+import { Modal } from '@/components/Modal';
+import Link from 'next/link';
+import { TextInput } from '@/components/formInputs';
+import { useRouter } from 'next/router';
+import DatasetImageCard from '@/components/DatasetImageCard';
+
+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 [newDatasetName, setNewDatasetName] = useState('');
+  const [isNewDatasetModalOpen, setIsNewDatasetModalOpen] = useState(false);
+  const [status, setStatus] = useState<'idle' | 'loading' | 'success' | 'error'>('idle');
+
+  const refreshImageList = (dbName: string) => {
+    setStatus('loading');
+    fetch('/api/datasets/listImages', {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+      },
+      body: JSON.stringify({ datasetName: dbName }),
+    })
+      .then(res => res.json())
+      .then(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) {
+      refreshImageList(datasetName);
+    }
+  }, [datasetName]);
+  return (
+    <>
+      <div className="space-y-6">
+        <div className="flex justify-between items-center">
+          <div>
+            <h1 className="text-xl font-bold mb-8">Dataset: {datasetName}</h1>
+          </div>
+        </div>
+        <Card title={`Images (${imgList.length})`}>
+          {status === 'loading' && <p>Loading...</p>}
+          {status === 'error' && <p>Error fetching images</p>}
+          {status === 'success' && (
+            <div className="grid grid-cols-1 sm:grid-cols-2 md:grid-cols-3 lg:grid-cols-4 gap-4">
+              {imgList.length === 0 && <p>No images found</p>}
+              {imgList.map(img => (
+                <DatasetImageCard
+                  key={img.img_path}
+                  alt="image"
+                  imageUrl={img.img_path}
+                />
+              ))}
+            </div>
+          )}
+        </Card>
+      </div>
+    </>
+  );
+}

ui/src/app/datasets/page.tsx

@@ -0,0 +1,124 @@
+'use client';
+
+import { useEffect, useState } from 'react';
+import Card from '@/components/Card';
+import { Modal } from '@/components/Modal';
+import Link from 'next/link';
+import { TextInput } from '@/components/formInputs';
+
+export default function Datasets() {
+  const [datasets, setDatasets] = useState([]);
+  const [newDatasetName, setNewDatasetName] = useState('');
+  const [isNewDatasetModalOpen, setIsNewDatasetModalOpen] = useState(false);
+  const [status, setStatus] = useState<'idle' | 'loading' | 'success' | 'error'>('idle');
+
+  const refreshDatasets = () => {
+    setStatus('loading');
+    fetch('/api/datasets/list')
+      .then(res => res.json())
+      .then(data => {
+        console.log('Datasets:', data);
+        // sort
+        data.sort((a: string, b: string) => a.localeCompare(b));
+        setDatasets(data);
+        setStatus('success');
+      })
+      .catch(error => {
+        console.error('Error fetching datasets:', error);
+        setStatus('error');
+      });
+  };
+  useEffect(() => {
+    refreshDatasets();
+  }, []);
+  return (
+    <>
+      <div className="space-y-6">
+        <div className="flex justify-between items-center">
+          <div>
+            <h1 className="text-xl font-bold mb-8">Datasets</h1>
+          </div>
+          <div>
+            <button
+              onClick={() => {
+                setIsNewDatasetModalOpen(true);
+              }}
+              className="bg-blue-500 hover:bg-blue-700 text-white font-bold py-2 px-4 rounded"
+            >
+              New Dataset
+            </button>
+          </div>
+        </div>
+        <Card title={`Datasets (${datasets.length})`}>
+          {status === 'loading' && <p>Loading...</p>}
+          {status === 'error' && <p>Error fetching datasets</p>}
+          {status === 'success' && (
+            <div className="space-y-1">
+              {datasets.length === 0 && <p>No datasets found</p>}
+              {datasets.map((dataset: string) => (
+                <Link href={`/datasets/${dataset}`} className="bg-gray-800 hover:bg-gray-700 py-2 px-4 rounded-lg cursor-pointer block" key={dataset}>
+                  {dataset}
+                </Link>
+              ))}
+            </div>
+          )}
+        </Card>
+      </div>
+      <Modal
+        isOpen={isNewDatasetModalOpen}
+        onClose={() => setIsNewDatasetModalOpen(false)}
+        title="New Dataset"
+        size="md"
+      >
+        <div className="space-y-4 text-gray-200">
+          <form
+            onSubmit={e => {
+              e.preventDefault();
+              console.log('Creating new dataset');
+              // make post with name to create new dataset
+              fetch('/api/datasets/create', {
+                method: 'POST',
+                headers: {
+                  'Content-Type': 'application/json',
+                },
+                body: JSON.stringify({ name: newDatasetName }),
+              })
+                .then(res => res.json())
+                .then(data => {
+                  console.log('New dataset created:', data);
+                  refreshDatasets();
+                  setNewDatasetName('');
+                  setIsNewDatasetModalOpen(false);
+                })
+                .catch(error => {
+                  console.error('Error creating new dataset:', error);
+                });
+            }}
+          >
+            <div className="text-sm text-gray-600">
+              This will create a new folder with the name below in your dataset folder.
+            </div>
+            <div>
+              <TextInput label="Dataset Name" value={newDatasetName} onChange={value => setNewDatasetName(value)} />
+            </div>
+
+            <div className="mt-6 flex justify-end space-x-3">
+              <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
+                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"
+                type="submit"
+              >
+                Confirm
+              </button>
+            </div>
+          </form>
+        </div>
+      </Modal>
+    </>
+  );
+}

ui/src/app/settings/page.tsx

@@ -6,6 +6,7 @@ export default function Settings() {
   const [settings, setSettings] = useState({
     HF_TOKEN: '',
     TRAINING_FOLDER: '',
+    DATASETS_FOLDER: '',
   });
   const [status, setStatus] = useState<'idle' | 'saving' | 'success' | 'error'>('idle');
 
@@ -17,6 +18,7 @@ export default function Settings() {
         setSettings({
           HF_TOKEN: data.HF_TOKEN || '',
           TRAINING_FOLDER: data.TRAINING_FOLDER || '',
+          DATASETS_FOLDER: data.DATASETS_FOLDER || '',
         });
       })
       .catch(error => console.error('Error fetching settings:', error));
@@ -65,7 +67,8 @@ export default function Settings() {
                 <a href="https://huggingface.co/settings/tokens" target="_blank" rel="noreferrer">
                   {' '}
                   Huggingface
-                </a> if you need to access gated/private models.
+                </a>{' '}
+                if you need to access gated/private models.
               </div>
             </label>
             <input
@@ -83,7 +86,8 @@ export default function Settings() {
             <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.
+                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
@@ -96,6 +100,28 @@ export default function Settings() {
               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>
 
         <button

ui/src/components/DatasetImageCard.tsx

@@ -0,0 +1,95 @@
+import React, { useRef, useEffect, useState, ReactNode } from 'react';
+
+interface DatasetImageCardProps {
+  imageUrl: string;
+  alt: string;
+  children?: ReactNode;
+  className?: string;
+}
+
+const DatasetImageCard: React.FC<DatasetImageCardProps> = ({ imageUrl, alt, children, className = '' }) => {
+  const cardRef = useRef<HTMLDivElement>(null);
+  const [isVisible, setIsVisible] = useState<boolean>(false);
+  const [loaded, setLoaded] = useState<boolean>(false);
+  const [isCaptionLoaded, setIsCaptionLoaded] = useState<boolean>(false);
+  const [caption, setCaption] = useState<string>('');
+  const isGettingCaption = useRef<boolean>(false);
+
+  const fetchCaption = async () => {
+    try {
+      if (isGettingCaption.current || isCaptionLoaded) return;
+      isGettingCaption.current = true;
+      const response = await fetch(`/api/caption/${encodeURIComponent(imageUrl)}`);
+      const data = await response.text();
+      setCaption(data);
+      setIsCaptionLoaded(true);
+    } catch (error) {
+      console.error('Error fetching caption:', error);
+    }
+  };
+
+  useEffect(() => {
+    isVisible && fetchCaption();
+  }, [isVisible]);
+
+  useEffect(() => {
+    // Create intersection observer to check visibility
+    const observer = new IntersectionObserver(
+      entries => {
+        if (entries[0].isIntersecting) {
+          setIsVisible(true);
+          observer.disconnect();
+        }
+      },
+      { threshold: 0.1 },
+    );
+
+    if (cardRef.current) {
+      observer.observe(cardRef.current);
+    }
+
+    return () => {
+      observer.disconnect();
+    };
+  }, []);
+
+  const handleLoad = (): void => {
+    setLoaded(true);
+  };
+
+  return (
+    <div className={`flex flex-col ${className}`}>
+      {/* Square image container */}
+      <div
+        ref={cardRef}
+        className="relative w-full"
+        style={{ paddingBottom: '100%' }} // Make it square
+      >
+        <div className="absolute inset-0 overflow-hidden rounded-t-lg shadow-md">
+          {isVisible && (
+            <img
+              src={`/api/img/${encodeURIComponent(imageUrl)}`}
+              alt={alt}
+              onLoad={handleLoad}
+              className={`w-full h-full object-contain transition-opacity duration-300 ${
+                loaded ? 'opacity-100' : 'opacity-0'
+              }`}
+            />
+          )}
+          {children && <div className="absolute inset-0 flex items-center justify-center">{children}</div>}
+        </div>
+      </div>
+
+      {/* Text area below the image */}
+      <div className="w-full p-2 bg-gray-800 text-white text-sm rounded-b-lg h-[75px]">
+        {isVisible && isCaptionLoaded && (
+          <form>
+            <textarea className="w-full bg-transparent resize-none" defaultValue={caption} rows={3} />
+          </form>
+        )}
+      </div>
+    </div>
+  );
+};
+
+export default DatasetImageCard;

ui/src/components/GPUMonitor.tsx

@@ -1,6 +1,7 @@
 // components/GpuMonitor.tsx
 import React, { useState, useEffect } from 'react';
 import { GPUApiResponse } from '@/types';
+import Loading from '@/components/Loading';
 
 const GpuMonitor: React.FC = () => {
   const [gpuData, setGpuData] = useState<GPUApiResponse | null>(null);
@@ -68,11 +69,7 @@ const GpuMonitor: React.FC = () => {
   };
 
   if (loading) {
-    return (
+    return <Loading />;
-      <div className="flex justify-center items-center h-64">
-        <div className="animate-spin rounded-full h-12 w-12 border-t-2 border-b-2 border-blue-500"></div>
-      </div>
-    );
   }
 
   if (error) {

ui/src/components/Loading.tsx

@@ -0,0 +1,7 @@
+export default function Loading() {
+  return (
+    <div className="flex justify-center items-center h-64">
+      <div className="animate-spin rounded-full h-12 w-12 border-t-2 border-b-2 border-blue-500"></div>
+    </div>
+  );
+}

ui/src/components/Modal.tsx

@@ -0,0 +1,110 @@
+import React, { Fragment, useEffect } from 'react';
+
+interface ModalProps {
+  isOpen: boolean;
+  onClose: () => void;
+  title?: string;
+  children: React.ReactNode;
+  showCloseButton?: boolean;
+  size?: 'sm' | 'md' | 'lg' | 'xl';
+  closeOnOverlayClick?: boolean;
+}
+
+export const Modal: React.FC<ModalProps> = ({
+  isOpen,
+  onClose,
+  title,
+  children,
+  showCloseButton = true,
+  size = 'md',
+  closeOnOverlayClick = true,
+}) => {
+  // Close on ESC key press
+  useEffect(() => {
+    const handleEscKey = (e: KeyboardEvent) => {
+      if (e.key === 'Escape' && isOpen) {
+        onClose();
+      }
+    };
+
+    if (isOpen) {
+      document.addEventListener('keydown', handleEscKey);
+      // Prevent body scrolling when modal is open
+      document.body.style.overflow = 'hidden';
+    }
+
+    return () => {
+      document.removeEventListener('keydown', handleEscKey);
+      document.body.style.overflow = 'auto';
+    };
+  }, [isOpen, onClose]);
+
+  // Handle overlay click
+  const handleOverlayClick = (e: React.MouseEvent<HTMLDivElement>) => {
+    if (e.target === e.currentTarget && closeOnOverlayClick) {
+      onClose();
+    }
+  };
+
+  if (!isOpen) return null;
+
+  // Size mapping
+  const sizeClasses = {
+    sm: 'max-w-md',
+    md: 'max-w-lg',
+    lg: 'max-w-2xl',
+    xl: 'max-w-4xl',
+  };
+
+  return (
+    <Fragment>
+      {/* Modal backdrop */}
+      <div
+        className="fixed inset-0 z-50 flex items-center justify-center overflow-y-auto bg-gray-900 bg-opacity-75 backdrop-blur-sm transition-opacity"
+        onClick={handleOverlayClick}
+        aria-modal="true"
+        role="dialog"
+        aria-labelledby="modal-title"
+      >
+        {/* Modal panel */}
+        <div
+          className={`relative mx-auto w-full ${sizeClasses[size]} rounded-lg bg-gray-800 border border-gray-700 shadow-xl transition-all`}
+          onClick={e => e.stopPropagation()}
+        >
+          {/* Modal header */}
+          {(title || showCloseButton) && (
+            <div className="flex items-center justify-between rounded-t-lg border-b border-gray-700 bg-gray-850 px-6 py-4">
+              {title && (
+                <h3 id="modal-title" className="text-xl font-semibold text-gray-100">
+                  {title}
+                </h3>
+              )}
+
+              {showCloseButton && (
+                <button
+                  type="button"
+                  className="ml-auto inline-flex items-center justify-center rounded-md p-2 text-gray-400 hover:bg-gray-700 hover:text-gray-200 focus:outline-none focus:ring-2 focus:ring-blue-500"
+                  onClick={onClose}
+                  aria-label="Close modal"
+                >
+                  <svg
+                    className="h-5 w-5"
+                    xmlns="http://www.w3.org/2000/svg"
+                    fill="none"
+                    viewBox="0 0 24 24"
+                    stroke="currentColor"
+                  >
+                    <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M6 18L18 6M6 6l12 12" />
+                  </svg>
+                </button>
+              )}
+            </div>
+          )}
+
+          {/* Modal content */}
+          <div className="px-6 py-4">{children}</div>
+        </div>
+      </div>
+    </Fragment>
+  );
+};

ui/src/components/Sidebar.tsx

@@ -1,10 +1,11 @@
 import Link from 'next/link';
-import { Home, Settings, BarChart2, BrainCircuit } from 'lucide-react';
+import { Home, Settings, BrainCircuit, Images } from 'lucide-react';
 
 const Sidebar = () => {
   const navigation = [
     { name: 'Dashboard', href: '/dashboard', icon: Home },
     { name: 'Train', href: '/train', icon: BrainCircuit },
+    { name: 'Datasets', href: '/datasets', icon: Images },
     { name: 'Settings', href: '/settings', icon: Settings },
   ];
 
@@ -12,8 +13,8 @@ const Sidebar = () => {
     <div className="flex flex-col w-64 bg-gray-900 text-gray-100">
       <div className="p-4">
         <h1 className="text-xl">
-            <img src="/ostris_logo.png" alt="Ostris AI Toolkit" className="w-auto h-8 mr-3 inline" />
+          <img src="/ostris_logo.png" alt="Ostris AI Toolkit" className="w-auto h-8 mr-3 inline" />
-            Ostris - AI Toolkit
+          Ostris - AI Toolkit
         </h1>
       </div>
       <nav className="flex-1">
@@ -31,6 +32,32 @@ const Sidebar = () => {
           ))}
         </ul>
       </nav>
+      <a href="https://patreon.com/ostris" target="_blank" rel="noreferrer" className="flex items-center space-x-2 p-4">
+      <div className='min-w-[26px] min-h-[26px]'>
+          <svg
+            viewBox="0 0 512 512"
+            xmlns="http://www.w3.org/2000/svg"
+            fillRule="evenodd"
+            clipRule="evenodd"
+            strokeLinejoin="round"
+            strokeMiterlimit="2"
+          >
+            <g transform="matrix(.47407 0 0 .47407 .383 .422)">
+              <clipPath id="prefix__a">
+                <path d="M0 0h1080v1080H0z"></path>
+              </clipPath>
+              <g clipPath="url(#prefix__a)">
+                <path
+                  d="M1033.05 324.45c-.19-137.9-107.59-250.92-233.6-291.7-156.48-50.64-362.86-43.3-512.28 27.2-181.1 85.46-237.99 272.66-240.11 459.36-1.74 153.5 13.58 557.79 241.62 560.67 169.44 2.15 194.67-216.18 273.07-321.33 55.78-74.81 127.6-95.94 216.01-117.82 151.95-37.61 255.51-157.53 255.29-316.38z"
+                  fillRule="nonzero"
+                  fill="#ffffff"
+                ></path>
+              </g>
+            </g>
+          </svg>
+        </div>
+        <div className="text-gray-500 text-md mb-2 flex-1 pt-2 pl-2">Support me on Patreon</div>
+      </a>
     </div>
   );
 };

ui/src/paths.ts

@@ -1,3 +1,4 @@
 import path from 'path';
 export const TOOLKIT_ROOT = path.resolve('@', '..', '..');
 export const defaultTrainFolder = path.join(TOOLKIT_ROOT, 'output');
+export const defaultDatasetsFolder = path.join(TOOLKIT_ROOT, 'datasets');