From d0214c0df9cbd5bca9265f1767095911fce1e3df Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Fri, 21 Feb 2025 06:18:27 -0700
Subject: [PATCH] Make ui more uniform

---
 ui/src/app/dashboard/page.tsx              |  16 +-
 ui/src/app/datasets/[datasetName]/page.tsx |   9 +-
 ui/src/app/datasets/page.tsx               |   9 +-
 ui/src/app/settings/page.tsx               | 169 +++--
 ui/src/app/train/page.tsx                  | 808 +++++++++++----------
 ui/src/components/layout.tsx               |  27 +
 6 files changed, 550 insertions(+), 488 deletions(-)
 create mode 100644 ui/src/components/layout.tsx

diff --git a/ui/src/app/dashboard/page.tsx b/ui/src/app/dashboard/page.tsx
index 94394919..2d989a18 100644
--- a/ui/src/app/dashboard/page.tsx
+++ b/ui/src/app/dashboard/page.tsx
@@ -1,12 +1,20 @@
 'use client';
 
 import GpuMonitor from '@/components/GPUMonitor';
+import { TopBar, MainContent } from '@/components/layout';
 
 export default function Dashboard() {
   return (
-    <div className="space-y-6">
-      <h1 className="text-xl font-bold mb-8">Dashboard</h1>
-      <GpuMonitor />
-    </div>
+    <>
+      <TopBar>
+        <div>
+          <h1 className="text-lg">Dashboard</h1>
+        </div>
+        <div className="flex-1"></div>
+      </TopBar>
+      <MainContent>
+        <GpuMonitor />
+      </MainContent>
+    </>
   );
 }
diff --git a/ui/src/app/datasets/[datasetName]/page.tsx b/ui/src/app/datasets/[datasetName]/page.tsx
index cb178279..37976431 100644
--- a/ui/src/app/datasets/[datasetName]/page.tsx
+++ b/ui/src/app/datasets/[datasetName]/page.tsx
@@ -5,6 +5,7 @@ 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';
 
 export default function DatasetPage({ params }: { params: { datasetName: string } }) {
   const [imgList, setImgList] = useState<{ img_path: string }[]>([]);
@@ -43,7 +44,7 @@ export default function DatasetPage({ params }: { params: { datasetName: string
   return (
     <>
       {/* Fixed top bar */}
-      <div className="absolute top-0 left-0 w-full h-12  dark:bg-gray-900 shadow-sm z-10 flex items-center px-2">
+      <TopBar>
         <div>
           <Button className="text-gray-500 dark:text-gray-300 px-3 mt-1" onClick={() => history.back()}>
             <FaChevronLeft />
@@ -61,8 +62,8 @@ export default function DatasetPage({ params }: { params: { datasetName: string
             Add Images
           </Button>
         </div>
-      </div>
-      <div className="pt-16 px-6 absolute top-0 left-0 w-full h-full overflow-auto">
+      </TopBar>
+      <MainContent>
         {status === 'loading' && <p>Loading...</p>}
         {status === 'error' && <p>Error fetching images</p>}
         {status === 'success' && (
@@ -78,7 +79,7 @@ export default function DatasetPage({ params }: { params: { datasetName: string
             ))}
           </div>
         )}
-      </div>
+      </MainContent>
       <AddImagesModal />
     </>
   );
diff --git a/ui/src/app/datasets/page.tsx b/ui/src/app/datasets/page.tsx
index 77851c18..2e2e43b3 100644
--- a/ui/src/app/datasets/page.tsx
+++ b/ui/src/app/datasets/page.tsx
@@ -8,6 +8,7 @@ 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';
 
 export default function Datasets() {
   const { datasets, status, refreshDatasets } = useDatasetList();
@@ -15,7 +16,7 @@ export default function Datasets() {
   const [isNewDatasetModalOpen, setIsNewDatasetModalOpen] = useState(false);
   return (
     <>
-      <div className="absolute top-0 left-0 w-full h-12  dark:bg-gray-900 shadow-sm z-10 flex items-center px-2">
+      <TopBar>
         <div>
           <h1 className="text-lg">Datasets</h1>
         </div>
@@ -28,8 +29,8 @@ export default function Datasets() {
             New Dataset
           </Button>
         </div>
-      </div>
-      <div className="pt-16 px-4 absolute top-0 left-0 w-full h-full overflow-auto">
+      </TopBar>
+      <MainContent>
         {status === 'loading' && <p>Loading...</p>}
         {status === 'error' && <p>Error fetching datasets</p>}
         {status === 'success' && (
@@ -79,7 +80,7 @@ export default function Datasets() {
             ))}
           </div>
         )}
-      </div>
+      </MainContent>
       <Modal
         isOpen={isNewDatasetModalOpen}
         onClose={() => setIsNewDatasetModalOpen(false)}
diff --git a/ui/src/app/settings/page.tsx b/ui/src/app/settings/page.tsx
index 3b50a282..0c8cdf6b 100644
--- a/ui/src/app/settings/page.tsx
+++ b/ui/src/app/settings/page.tsx
@@ -2,6 +2,7 @@
 
 import { useEffect, useState } from 'react';
 import useSettings from '@/hooks/useSettings';
+import { TopBar, MainContent } from '@/components/layout';
 
 export default function Settings() {
   const { settings, setSettings } = useSettings();
@@ -37,87 +38,97 @@ export default function Settings() {
   };
 
   return (
-    <div className="max-w-2xl mx-auto space-y-6">
-      <h1 className="text-3xl font-bold mb-8">Settings</h1>
-
-      <form onSubmit={handleSubmit} className="space-y-6">
-        <div className="space-y-4">
-          <div>
-            <label htmlFor="HF_TOKEN" className="block text-sm font-medium mb-2">
-              Hugging Face Token
-              <div className="text-gray-500 text-sm ml-1">
-                Create a Read token on{' '}
-                <a href="https://huggingface.co/settings/tokens" target="_blank" rel="noreferrer">
-                  {' '}
-                  Huggingface
-                </a>{' '}
-                if you need to access gated/private models.
-              </div>
-            </label>
-            <input
-              type="password"
-              id="HF_TOKEN"
-              name="HF_TOKEN"
-              value={settings.HF_TOKEN}
-              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 your Hugging Face token"
-            />
-          </div>
-
-          <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>
+    <>
+      <TopBar>
+        <div>
+          <h1 className="text-lg">Settings</h1>
         </div>
+        <div className="flex-1"></div>
+      </TopBar>
+      <MainContent>
+        <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="HF_TOKEN" className="block text-sm font-medium mb-2">
+                    Hugging Face Token
+                    <div className="text-gray-500 text-sm ml-1">
+                      Create a Read token on{' '}
+                      <a href="https://huggingface.co/settings/tokens" target="_blank" rel="noreferrer">
+                        {' '}
+                        Huggingface
+                      </a>{' '}
+                      if you need to access gated/private models.
+                    </div>
+                  </label>
+                  <input
+                    type="password"
+                    id="HF_TOKEN"
+                    name="HF_TOKEN"
+                    value={settings.HF_TOKEN}
+                    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 your Hugging Face token"
+                  />
+                </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>
+                <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>
 
-        {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>
-    </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>
+
+          <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>
+    </>
   );
 }
diff --git a/ui/src/app/train/page.tsx b/ui/src/app/train/page.tsx
index f7936b6a..0938308b 100644
--- a/ui/src/app/train/page.tsx
+++ b/ui/src/app/train/page.tsx
@@ -1,7 +1,6 @@
 'use client';
-// todo update training folder from settings
 
-import { use, useEffect, useMemo, useState } from 'react';
+import { useEffect, useState } from 'react';
 import { useSearchParams, useRouter } from 'next/navigation';
 import { options } from './options';
 import { defaultJobConfig, defaultDatasetConfig } from './jobConfig';
@@ -15,6 +14,7 @@ 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';
 
 export default function TrainingForm() {
   const router = useRouter();
@@ -37,7 +37,7 @@ export default function TrainingForm() {
     setDatasetOptions(datasetOptions);
     const defaultDatasetPath = defaultDatasetConfig.folder_path;
 
-    for(let i = 0; i < jobConfig.config.process[0].datasets.length; i++) {
+    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) {
         setJobConfig(datasetOptions[0].value, `config.process[0].datasets[${i}].folder_path`);
@@ -105,218 +105,223 @@ export default function TrainingForm() {
   };
 
   return (
-    <div className="space-y-6">
-      <h1 className="text-xl font-bold mb-8">{runId ? 'Edit Training Run' : 'New Training Run'}</h1>
-
-      <form onSubmit={handleSubmit} className="space-y-8">
-        <div className="grid grid-cols-1 md:grid-cols-2 lg:grid-cols-4 gap-6">
-          <Card title="Job Settings">
-            <TextInput
-              label="Training Name"
-              value={jobConfig.config.name}
-              onChange={value => setJobConfig(value, 'config.name')}
-              placeholder="Enter training name"
-              required
-            />
-            <SelectInput
-              label="GPU ID"
-              value={`${gpuID}`}
-              className="pt-2"
-              onChange={value => setGpuID(parseInt(value))}
-              options={gpuList.map(gpu => ({ value: `${gpu}`, label: `GPU #${gpu}` }))}
-            />
-          </Card>
-
-          {/* Model Configuration Section */}
-          <Card title="Model Configuration">
-            <SelectInput
-              label="Name or Path"
-              value={jobConfig.config.process[0].model.name_or_path}
-              onChange={value => {
-                // see if model changed
-                const currentModel = options.model.find(
-                  model => model.name_or_path === jobConfig.config.process[0].model.name_or_path,
-                );
-                if (!currentModel || currentModel.name_or_path === value) {
-                  // model has not changed
-                  return;
-                }
-                // revert defaults from previous model
-                for (const key in currentModel.defaults) {
-                  setJobConfig(currentModel.defaults[key][1], key);
-                }
-                // set new model
-                setJobConfig(value, 'config.process[0].model.name_or_path');
-                // update the defaults when a model is selected
-                const model = options.model.find(model => model.name_or_path === value);
-                if (model?.defaults) {
-                  for (const key in model.defaults) {
-                    setJobConfig(model.defaults[key][0], key);
-                  }
-                }
-              }}
-              options={options.model.map(model => ({
-                value: model.name_or_path,
-                label: model.name_or_path,
-              }))}
-            />
-            <FormGroup label="Quantize" className="pt-2">
-              <Checkbox
-                label="Transformer"
-                checked={jobConfig.config.process[0].model.quantize}
-                onChange={value => setJobConfig(value, 'config.process[0].model.quantize')}
+    <>
+      <TopBar>
+        <div>
+          <h1 className="text-lg">{runId ? 'Edit Training Run' : 'New Training Run'}</h1>
+        </div>
+        <div className="flex-1"></div>
+      </TopBar>
+      <MainContent>
+        <form onSubmit={handleSubmit} className="space-y-8">
+          <div className="grid grid-cols-1 md:grid-cols-2 lg:grid-cols-4 gap-6">
+            <Card title="Job Settings">
+              <TextInput
+                label="Training Name"
+                value={jobConfig.config.name}
+                onChange={value => setJobConfig(value, 'config.name')}
+                placeholder="Enter training name"
+                required
               />
-              <Checkbox
-                label="Text Encoder"
-                checked={jobConfig.config.process[0].model.quantize_te}
-                onChange={value => setJobConfig(value, 'config.process[0].model.quantize_te')}
+              <SelectInput
+                label="GPU ID"
+                value={`${gpuID}`}
+                className="pt-2"
+                onChange={value => setGpuID(parseInt(value))}
+                options={gpuList.map(gpu => ({ value: `${gpu}`, label: `GPU #${gpu}` }))}
               />
-            </FormGroup>
-          </Card>
-          {jobConfig.config.process[0].network?.linear && (
-            <Card title="LoRA Configuration">
-              <NumberInput
-                label="Linear Rank"
-                value={jobConfig.config.process[0].network.linear}
+            </Card>
+
+            {/* Model Configuration Section */}
+            <Card title="Model Configuration">
+              <SelectInput
+                label="Name or Path"
+                value={jobConfig.config.process[0].model.name_or_path}
                 onChange={value => {
-                  setJobConfig(value, 'config.process[0].network.linear');
-                  setJobConfig(value, 'config.process[0].network.linear_alpha');
+                  // see if model changed
+                  const currentModel = options.model.find(
+                    model => model.name_or_path === jobConfig.config.process[0].model.name_or_path,
+                  );
+                  if (!currentModel || currentModel.name_or_path === value) {
+                    // model has not changed
+                    return;
+                  }
+                  // revert defaults from previous model
+                  for (const key in currentModel.defaults) {
+                    setJobConfig(currentModel.defaults[key][1], key);
+                  }
+                  // set new model
+                  setJobConfig(value, 'config.process[0].model.name_or_path');
+                  // update the defaults when a model is selected
+                  const model = options.model.find(model => model.name_or_path === value);
+                  if (model?.defaults) {
+                    for (const key in model.defaults) {
+                      setJobConfig(model.defaults[key][0], key);
+                    }
+                  }
                 }}
-                placeholder="eg. 16"
+                options={options.model.map(model => ({
+                  value: model.name_or_path,
+                  label: model.name_or_path,
+                }))}
+              />
+              <FormGroup label="Quantize" className="pt-2">
+                <Checkbox
+                  label="Transformer"
+                  checked={jobConfig.config.process[0].model.quantize}
+                  onChange={value => setJobConfig(value, 'config.process[0].model.quantize')}
+                />
+                <Checkbox
+                  label="Text Encoder"
+                  checked={jobConfig.config.process[0].model.quantize_te}
+                  onChange={value => setJobConfig(value, 'config.process[0].model.quantize_te')}
+                />
+              </FormGroup>
+            </Card>
+            {jobConfig.config.process[0].network?.linear && (
+              <Card title="LoRA Configuration">
+                <NumberInput
+                  label="Linear Rank"
+                  value={jobConfig.config.process[0].network.linear}
+                  onChange={value => {
+                    setJobConfig(value, 'config.process[0].network.linear');
+                    setJobConfig(value, 'config.process[0].network.linear_alpha');
+                  }}
+                  placeholder="eg. 16"
+                  min={1}
+                  max={1024}
+                  required
+                />
+              </Card>
+            )}
+            <Card title="Save Configuration">
+              <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}
-                max={1024}
                 required
               />
             </Card>
-          )}
-          <Card title="Save Configuration">
-            <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 Configuration">
-            <div className="grid grid-cols-1 md:grid-cols-2 lg:grid-cols-4 gap-6">
-              <div>
-                <NumberInput
-                  label="Batch Size"
-                  className="pt-2"
-                  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>
+            <Card title="Training Configuration">
+              <div className="grid grid-cols-1 md:grid-cols-2 lg:grid-cols-4 gap-6">
+                <div>
+                  <NumberInput
+                    label="Batch Size"
+                    className="pt-2"
+                    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"
+                    className="pt-2"
+                    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>
-              <div>
-                <SelectInput
-                  label="Optimizer"
-                  className="pt-2"
-                  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>
-          </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}
-                      />
-                      {/* <TextInput
+            </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}
+                        />
+                        {/* <TextInput
                         label="Folder Path"
                         value={dataset.folder_path}
                         onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].folder_path`)}
                         placeholder="eg. /path/to/images/folder"
                         required
                       /> */}
-                      {/* <TextInput
+                        {/* <TextInput
                         label="Mask Folder Path"
                         className="pt-2"
                         value={dataset.mask_path || ''}
@@ -339,211 +344,220 @@ export default function TrainingForm() {
                         max={1}
                         required
                       /> */}
-                    </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"
-                      />
-                      <TextInput
-                        label="Caption Extension"
-                        className="pt-2"
-                        value={dataset.caption_ext}
-                        onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].caption_ext`)}
-                        placeholder="eg. txt"
-                        required
-                      />
-                      <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
-                      />
-                    </div>
-                    <div>
-                      <FormGroup label="Settings" className="">
-                        <Checkbox
-                          label="Cache Latents to Disk"
-                          checked={dataset.cache_latents_to_disk || false}
-                          onChange={value =>
-                            setJobConfig(value, `config.process[0].datasets[${i}].cache_latents_to_disk`)
-                          }
+                      </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"
                         />
-                        <Checkbox
-                          label="Is Regularization"
+                        <TextInput
+                          label="Caption Extension"
                           className="pt-2"
-                          checked={dataset.is_reg || false}
-                          onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].is_reg`)}
+                          value={dataset.caption_ext}
+                          onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].caption_ext`)}
+                          placeholder="eg. txt"
+                          required
                         />
-                      </FormGroup>
+                        <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
+                        />
+                      </div>
+                      <div>
+                        <FormGroup label="Settings" className="">
+                          <Checkbox
+                            label="Cache Latents to Disk"
+                            checked={dataset.cache_latents_to_disk || false}
+                            onChange={value =>
+                              setJobConfig(value, `config.process[0].datasets[${i}].cache_latents_to_disk`)
+                            }
+                          />
+                          <Checkbox
+                            label="Is Regularization"
+                            className="pt-2"
+                            checked={dataset.is_reg || false}
+                            onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].is_reg`)}
+                          />
+                        </FormGroup>
+                      </div>
+                      <div>
+                        <FormGroup label="Resolutions" className="pt-2">
+                          {[256, 512, 768, 1024, 1280].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`);
+                              }}
+                            />
+                          ))}
+                        </FormGroup>
+                      </div>
+                    </div>
+                  </div>
+                ))}
+                <button
+                  type="button"
+                  onClick={() =>
+                    setJobConfig(
+                      [...jobConfig.config.process[0].datasets, objectCopy(defaultDatasetConfig)],
+                      '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 Configuration">
+              <div className="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' }]}
+                  />
+                </div>
+                <div>
+                  <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"
+                    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={256}
+                    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={256}
+                    required
+                  />
+                </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={`Sample Prompts (${jobConfig.config.process[0].sample.prompts.length})`}
+                className="pt-2"
+              >
+                {jobConfig.config.process[0].sample.prompts.map((prompt, i) => (
+                  <div key={i} className="flex items-center space-x-2">
+                    <div className="flex-1">
+                      <TextInput
+                        value={prompt}
+                        onChange={value => setJobConfig(value, `config.process[0].sample.prompts[${i}]`)}
+                        placeholder="Enter prompt"
+                        required
+                      />
                     </div>
                     <div>
-                      <FormGroup label="Resolutions" className="pt-2">
-                        {[256, 512, 768, 1024, 1280].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`);
-                            }}
-                          />
-                        ))}
-                      </FormGroup>
+                      <button
+                        type="button"
+                        onClick={() =>
+                          setJobConfig(
+                            jobConfig.config.process[0].sample.prompts.filter((_, index) => index !== i),
+                            'config.process[0].sample.prompts',
+                          )
+                        }
+                        className="rounded-full p-1 text-sm"
+                      >
+                        <X />
+                      </button>
                     </div>
                   </div>
-                </div>
-              ))}
-              <button
-                type="button"
-                onClick={() =>
-                  setJobConfig(
-                    [...jobConfig.config.process[0].datasets, objectCopy(defaultDatasetConfig)],
-                    '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 Configuration">
-            <div className="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' }]}
-                />
-              </div>
-              <div>
-                <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"
-                  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={256}
-                  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={256}
-                  required
-                />
-              </div>
+                ))}
+                <button
+                  type="button"
+                  onClick={() =>
+                    setJobConfig(
+                      [...jobConfig.config.process[0].sample.prompts, ''],
+                      'config.process[0].sample.prompts',
+                    )
+                  }
+                  className="w-full px-4 py-2 bg-gray-700 hover:bg-gray-600 rounded-lg transition-colors"
+                >
+                  Add Prompt
+                </button>
+              </FormGroup>
+            </Card>
+          </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={`Sample Prompts (${jobConfig.config.process[0].sample.prompts.length})`} className="pt-2">
-              {jobConfig.config.process[0].sample.prompts.map((prompt, i) => (
-                <div key={i} className="flex items-center space-x-2">
-                  <div className="flex-1">
-                    <TextInput
-                      value={prompt}
-                      onChange={value => setJobConfig(value, `config.process[0].sample.prompts[${i}]`)}
-                      placeholder="Enter prompt"
-                      required
-                    />
-                  </div>
-                  <div>
-                    <button
-                      type="button"
-                      onClick={() =>
-                        setJobConfig(
-                          jobConfig.config.process[0].sample.prompts.filter((_, index) => index !== i),
-                          'config.process[0].sample.prompts',
-                        )
-                      }
-                      className="rounded-full p-1 text-sm"
-                    >
-                      <X />
-                    </button>
-                  </div>
-                </div>
-              ))}
-              <button
-                type="button"
-                onClick={() =>
-                  setJobConfig([...jobConfig.config.process[0].sample.prompts, ''], 'config.process[0].sample.prompts')
-                }
-                className="w-full px-4 py-2 bg-gray-700 hover:bg-gray-600 rounded-lg transition-colors"
-              >
-                Add Prompt
-              </button>
-            </FormGroup>
-          </Card>
-        </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...' : runId ? 'Update Training' : 'Create Training'}
+          </button>
 
-        <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...' : runId ? 'Update Training' : 'Create Training'}
-        </button>
-
-        {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>
-    </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>
+      </MainContent>
+    </>
   );
 }
diff --git a/ui/src/components/layout.tsx b/ui/src/components/layout.tsx
new file mode 100644
index 00000000..a23cc45b
--- /dev/null
+++ b/ui/src/components/layout.tsx
@@ -0,0 +1,27 @@
+import classNames from 'classnames';
+
+interface Props {
+  className?: string;
+  children?: React.ReactNode;
+}
+
+export const TopBar: React.FC<Props> = ({ children, className }) => {
+  return (
+    <div
+      className={classNames(
+        'absolute top-0 left-0 w-full h-12  dark:bg-gray-900 shadow-sm z-10 flex items-center px-2',
+        className,
+      )}
+    >
+      {children ? children : null}
+    </div>
+  );
+};
+
+export const MainContent: React.FC<Props> = ({ children, className }) => {
+  return (
+    <div className={classNames('pt-16 px-4 absolute top-0 left-0 w-full h-full overflow-auto', className)}>
+      {children ? children : null}
+    </div>
+  );
+};