github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Jaret Burkett
2025-03-20 08:07:09 -06:00
parent bbfd6ef0fe
commit 3a6b24f4c8
47 changed files with 618 additions and 378 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
ui/src/app/api/auth/route.ts Normal file
View File

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

2
ui/src/app/api/datasets/create/route.tsx
View File

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

28
ui/src/app/api/datasets/listImages/route.ts
View File

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

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

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

9
ui/src/app/api/gpu/route.ts
View File

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

1
ui/src/app/api/img/caption/route.ts
View File

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

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

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

2
ui/src/app/api/jobs/route.ts
View File

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

2
ui/src/app/api/settings/route.ts
View File

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

15
ui/src/app/datasets/[datasetName]/page.tsx
View File

@@ -6,6 +6,7 @@ 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';
export default function DatasetPage({ params }: { params: { datasetName: string } }) {
const [imgList, setImgList] = useState<{ img_path: string }[]>([]);
@@ -15,15 +16,11 @@ export default function DatasetPage({ params }: { params: { datasetName: string
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('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));

24
ui/src/app/datasets/page.tsx
View File

@@ -10,6 +10,7 @@ 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';
export default function Datasets() {
const { datasets, status, refreshDatasets } = useDatasetList();
@@ -54,16 +55,10 @@ export default function Datasets() {
type: 'warning',
confirmText: 'Delete',
onConfirm: () => {
fetch('/api/datasets/delete', {
method: 'POST',
headers: {
'Content-Type': 'application/json',
},
body: JSON.stringify({ name: datasetName }),
})
.then(res => res.json())
.then(data => {
console.log('Dataset deleted:', data);
apiClient
.post('/api/datasets/delete', { name: datasetName })
.then(() => {
console.log('Dataset deleted:', datasetName);
refreshDatasets();
})
.catch(error => {
@@ -76,14 +71,7 @@ export default function Datasets() {
const handleCreateDataset = async (e: React.FormEvent) => {
e.preventDefault();
try {
const response = await fetch('/api/datasets/create', {
method: 'POST',
headers: {
'Content-Type': 'application/json',
},
body: JSON.stringify({ name: newDatasetName }),
});
const data = await response.json();
const data = await apiClient.post('/api/datasets/create', { name: newDatasetName }).then(res => res.data);
console.log('New dataset created:', data);
refreshDatasets();
setNewDatasetName('');

6
ui/src/app/jobs/new/AdvancedJob.tsx
View File

@@ -33,11 +33,7 @@ const yamlConfig: YAML.DocumentOptions &
directives: true,
};
export default function AdvancedJob({
jobConfig,
setJobConfig,
settings,
}: Props) {
export default function AdvancedJob({ jobConfig, setJobConfig, settings }: Props) {
const [editorValue, setEditorValue] = useState<string>('');
const lastJobConfigUpdateStringRef = useRef('');
const editorRef = useRef<editor.IStandaloneCodeEditor | null>(null);

11
ui/src/app/jobs/new/jobConfig.ts
View File

@@ -30,7 +30,7 @@ export const defaultJobConfig: JobConfig = {
linear: 32,
linear_alpha: 32,
lokr_full_rank: true,
lokr_factor: -1
lokr_factor: -1,
},
save: {
dtype: 'bf16',
@@ -39,9 +39,7 @@ export const defaultJobConfig: JobConfig = {
save_format: 'diffusers',
push_to_hub: false,
},
datasets: [
defaultDatasetConfig
],
datasets: [defaultDatasetConfig],
train: {
batch_size: 1,
bypass_guidance_embedding: true,
@@ -55,7 +53,7 @@ export const defaultJobConfig: JobConfig = {
timestep_type: 'sigmoid',
content_or_style: 'balanced',
optimizer_params: {
weight_decay: 1e-4
weight_decay: 1e-4,
},
unload_text_encoder: false,
lr: 0.0001,
@@ -66,8 +64,7 @@ export const defaultJobConfig: JobConfig = {
dtype: 'bf16',
diff_output_preservation: false,
diff_output_preservation_multiplier: 1.0,
diff_output_preservation_class: 'person'
diff_output_preservation_class: 'person',
},
model: {
name_or_path: 'ostris/Flex.1-alpha',

2
ui/src/app/jobs/new/options.ts
View File

@@ -12,7 +12,7 @@ export const modelArchs = [
{ name: 'flux', label: 'Flux.1' },
{ name: 'wan21', label: 'Wan 2.1' },
{ name: 'lumina2', label: 'Lumina2' },
]
];
export const isVideoModelFromArch = (arch: string) => {
const videoArches = ['wan21'];

59
ui/src/app/jobs/new/page.tsx
View File

@@ -19,6 +19,7 @@ import { Button } from '@headlessui/react';
import { FaChevronLeft } from 'react-icons/fa';
import SimpleJob from './SimpleJob';
import AdvancedJob from './AdvancedJob';
import { apiClient } from '@/utils/api';
const isDev = process.env.NODE_ENV === 'development';
@@ -56,12 +57,13 @@ export default function TrainingForm() {
useEffect(() => {
if (runId) {
fetch(`/api/jobs?id=${runId}`)
.then(res => res.json())
apiClient
.get(`/api/jobs?id=${runId}`)
.then(res => res.data)
.then(data => {
console.log('Training:', data);
setGpuIDs(data.gpu_ids);
setJobConfig(JSON.parse(data.job_config));
// setJobConfig(data.name, 'config.name');
})
.catch(error => console.error('Error fetching training:', error));
}
@@ -85,33 +87,30 @@ export default function TrainingForm() {
if (status === 'saving') return;
setStatus('saving');
try {
const response = await fetch('/api/jobs', {
method: 'POST',
headers: {
'Content-Type': 'application/json',
},
body: JSON.stringify({
id: runId,
name: jobConfig.config.name,
gpu_ids: gpuIDs,
job_config: jobConfig,
}),
});
if (!response.ok) throw new Error('Failed to save training');
setStatus('success');
if (!runId) {
const data = await response.json();
router.push(`/jobs/${data.id}`);
}
setTimeout(() => setStatus('idle'), 2000);
} catch (error) {
console.error('Error saving training:', error);
setStatus('error');
setTimeout(() => setStatus('idle'), 2000);
}
apiClient
.post('/api/jobs', {
id: runId,
name: jobConfig.config.name,
gpu_ids: gpuIDs,
job_config: jobConfig,
})
.then(res => {
setStatus('success');
if (runId) {
router.push(`/jobs/${runId}`);
} else {
router.push(`/jobs/${res.data.id}`);
}
})
.catch(error => {
console.error('Error saving training:', error);
setStatus('error');
})
.finally(() =>
setTimeout(() => {
setStatus('idle');
}, 2000),
);
};
const handleSubmit = async (e: React.FormEvent) => {

5
ui/src/app/jobs/page.tsx
View File

@@ -13,10 +13,7 @@ export default function Dashboard() {
</div>
<div className="flex-1"></div>
<div>
<Link
href="/jobs/new"
className="text-gray-200 bg-slate-600 px-3 py-1 rounded-md"
>
<Link href="/jobs/new" className="text-gray-200 bg-slate-600 px-3 py-1 rounded-md">
New Training Job
</Link>
</div>

19
ui/src/app/layout.tsx
View File

@@ -6,6 +6,7 @@ import { ThemeProvider } from '@/components/ThemeProvider';
import ConfirmModal from '@/components/ConfirmModal';
import SampleImageModal from '@/components/SampleImageModal';
import { Suspense } from 'react';
import AuthWrapper from '@/components/AuthWrapper';
const inter = Inter({ subsets: ['latin'] });
@@ -15,6 +16,9 @@ export const metadata: Metadata = {
};
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>
@@ -22,13 +26,14 @@ export default function RootLayout({ children }: { children: React.ReactNode })
</head>
<body className={inter.className}>
<ThemeProvider>
<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 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>
</ThemeProvider>
<ConfirmModal />
<SampleImageModal />

2
ui/src/app/page.tsx
View File

@@ -2,4 +2,4 @@ import { redirect } from 'next/navigation';
export default function Home() {
redirect('/dashboard');
}
}

29
ui/src/app/settings/page.tsx
View File

@@ -3,6 +3,7 @@
import { useEffect, useState } from 'react';
import useSettings from '@/hooks/useSettings';
import { TopBar, MainContent } from '@/components/layout';
import { apiClient } from '@/utils/api';
export default function Settings() {
const { settings, setSettings } = useSettings();
@@ -12,24 +13,18 @@ export default function Settings() {
e.preventDefault();
setStatus('saving');
try {
const response = await fetch('/api/settings', {
method: 'POST',
headers: {
'Content-Type': 'application/json',
},
body: JSON.stringify(settings),
apiClient
.post('/api/settings', settings)
.then(() => {
setStatus('success');
})
.catch(error => {
console.error('Error saving settings:', error);
setStatus('error');
})
.finally(() => {
setTimeout(() => setStatus('idle'), 2000);
});
if (!response.ok) throw new Error('Failed to save settings');
setStatus('success');
setTimeout(() => setStatus('idle'), 2000);
} catch (error) {
console.error('Error saving settings:', error);
setStatus('error');
setTimeout(() => setStatus('idle'), 2000);
}
};
const handleChange = (e: React.ChangeEvent<HTMLInputElement>) => {