From 306a9ac32e7e93951fbff6b96820a32b372fe7cc Mon Sep 17 00:00:00 2001 From: Riffmaster2001 Date: Sat, 7 Feb 2026 14:16:31 -0800 Subject: [PATCH] adding a new feature for people running ai-toolkit on the same server as their inference tools like ComfyUI. The change allows you to put in a lora install path. If that path is not empty a new install icon appears next to completed loras next to the download button. This allows you to 1click install your loras to test them directly. --- ui/src/app/api/files/install/route.ts | 62 +++++++++++++++++++++++ ui/src/app/api/settings/route.ts | 9 +++- ui/src/app/settings/page.tsx | 18 +++++++ ui/src/components/FilesWidget.tsx | 72 +++++++++++++++++++++++---- ui/src/hooks/useSettings.tsx | 3 ++ ui/src/server/settings.ts | 19 +++++++ 6 files changed, 170 insertions(+), 13 deletions(-) create mode 100644 ui/src/app/api/files/install/route.ts diff --git a/ui/src/app/api/files/install/route.ts b/ui/src/app/api/files/install/route.ts new file mode 100644 index 000000000..8168bffb4 --- /dev/null +++ b/ui/src/app/api/files/install/route.ts @@ -0,0 +1,62 @@ +import { NextRequest, NextResponse } from 'next/server'; +import fs from 'fs'; +import path from 'path'; +import { getDatasetsRoot, getTrainingFolder, getLoraInstallPath } from '@/server/settings'; + +export async function POST(request: NextRequest) { + try { + const body = await request.json(); + const { filePath } = body; + + if (!filePath || typeof filePath !== 'string') { + return NextResponse.json({ error: 'filePath is required' }, { status: 400 }); + } + + const loraInstallPath = await getLoraInstallPath(); + if (!loraInstallPath) { + return NextResponse.json({ error: 'LoRA Install Path is not configured in settings' }, { status: 400 }); + } + + // Security check: ensure source path is in allowed directories + const datasetRoot = await getDatasetsRoot(); + const trainingRoot = await getTrainingFolder(); + const allowedDirs = [datasetRoot, trainingRoot]; + + const isAllowed = + allowedDirs.some(allowedDir => filePath.startsWith(allowedDir)) && !filePath.includes('..'); + + if (!isAllowed) { + return NextResponse.json({ error: 'Access denied' }, { status: 403 }); + } + + // Verify source file exists + if (!fs.existsSync(filePath)) { + return NextResponse.json({ error: 'Source file not found' }, { status: 404 }); + } + + const stat = fs.statSync(filePath); + if (!stat.isFile()) { + return NextResponse.json({ error: 'Source path is not a file' }, { status: 400 }); + } + + // Verify destination directory exists + if (!fs.existsSync(loraInstallPath)) { + return NextResponse.json({ error: `Install directory does not exist: ${loraInstallPath}` }, { status: 400 }); + } + + const destStat = fs.statSync(loraInstallPath); + if (!destStat.isDirectory()) { + return NextResponse.json({ error: 'Install path is not a directory' }, { status: 400 }); + } + + const fileName = path.basename(filePath); + const destPath = path.join(loraInstallPath, fileName); + + await fs.promises.copyFile(filePath, destPath); + + return NextResponse.json({ success: true, destPath }); + } catch (error) { + console.error('Error installing file:', error); + return NextResponse.json({ error: 'Failed to install file' }, { status: 500 }); + } +} diff --git a/ui/src/app/api/settings/route.ts b/ui/src/app/api/settings/route.ts index 62528cdd0..01ca598e2 100644 --- a/ui/src/app/api/settings/route.ts +++ b/ui/src/app/api/settings/route.ts @@ -29,9 +29,9 @@ export async function GET() { export async function POST(request: Request) { try { const body = await request.json(); - const { HF_TOKEN, TRAINING_FOLDER, DATASETS_FOLDER } = body; + const { HF_TOKEN, TRAINING_FOLDER, DATASETS_FOLDER, LORA_INSTALL_PATH } = body; - // Upsert both settings + // Upsert all settings await Promise.all([ prisma.settings.upsert({ where: { key: 'HF_TOKEN' }, @@ -48,6 +48,11 @@ export async function POST(request: Request) { update: { value: DATASETS_FOLDER }, create: { key: 'DATASETS_FOLDER', value: DATASETS_FOLDER }, }), + prisma.settings.upsert({ + where: { key: 'LORA_INSTALL_PATH' }, + update: { value: LORA_INSTALL_PATH || '' }, + create: { key: 'LORA_INSTALL_PATH', value: LORA_INSTALL_PATH || '' }, + }), ]); flushCache(); diff --git a/ui/src/app/settings/page.tsx b/ui/src/app/settings/page.tsx index 4bf257c31..16abc89c7 100644 --- a/ui/src/app/settings/page.tsx +++ b/ui/src/app/settings/page.tsx @@ -108,6 +108,24 @@ export default function Settings() { placeholder="Enter datasets folder path" /> +
+ + +
diff --git a/ui/src/components/FilesWidget.tsx b/ui/src/components/FilesWidget.tsx index b73dabf09..ebd552726 100644 --- a/ui/src/components/FilesWidget.tsx +++ b/ui/src/components/FilesWidget.tsx @@ -1,10 +1,42 @@ -import React from 'react'; +import React, { useEffect, useState } from 'react'; import useFilesList from '@/hooks/useFilesList'; -import Link from 'next/link'; -import { Loader2, AlertCircle, Download, Box, Brain } from 'lucide-react'; +import { Loader2, AlertCircle, Download, Box, Brain, FolderInput, Check } from 'lucide-react'; +import { apiClient } from '@/utils/api'; export default function FilesWidget({ jobID }: { jobID: string }) { const { files, status, refreshFiles } = useFilesList(jobID, 5000); + const [loraInstallPath, setLoraInstallPath] = useState(''); + const [installingFiles, setInstallingFiles] = useState>( + {}, + ); + + useEffect(() => { + apiClient + .get('/api/settings') + .then(res => res.data) + .then(data => { + setLoraInstallPath(data.LORA_INSTALL_PATH || ''); + }) + .catch(() => {}); + }, []); + + const handleInstall = async (filePath: string, e: React.MouseEvent) => { + e.preventDefault(); + e.stopPropagation(); + setInstallingFiles(prev => ({ ...prev, [filePath]: 'installing' })); + try { + await apiClient.post('/api/files/install', { filePath }); + setInstallingFiles(prev => ({ ...prev, [filePath]: 'success' })); + setTimeout(() => { + setInstallingFiles(prev => ({ ...prev, [filePath]: 'idle' })); + }, 2000); + } catch { + setInstallingFiles(prev => ({ ...prev, [filePath]: 'error' })); + setTimeout(() => { + setInstallingFiles(prev => ({ ...prev, [filePath]: 'idle' })); + }, 2000); + } + }; const cleanSize = (size: number) => { if (size < 1024) { @@ -47,11 +79,10 @@ export default function FilesWidget({ jobID }: { jobID: string }) { {files.map((file, index) => { const fileName = file.path.split('/').pop() || ''; const nameWithoutExt = fileName.replace('.safetensors', ''); + const installStatus = installingFiles[file.path] || 'idle'; return ( -
@@ -65,13 +96,32 @@ export default function FilesWidget({ jobID }: { jobID: string }) { .safetensors
-
- {cleanSize(file.size)} -
+
+ {cleanSize(file.size)} + {loraInstallPath && ( + + )} + -
+
- +
); })} diff --git a/ui/src/hooks/useSettings.tsx b/ui/src/hooks/useSettings.tsx index 35fcc5383..d3b8a00ee 100644 --- a/ui/src/hooks/useSettings.tsx +++ b/ui/src/hooks/useSettings.tsx @@ -7,6 +7,7 @@ export interface Settings { HF_TOKEN: string; TRAINING_FOLDER: string; DATASETS_FOLDER: string; + LORA_INSTALL_PATH: string; } export default function useSettings() { @@ -14,6 +15,7 @@ export default function useSettings() { HF_TOKEN: '', TRAINING_FOLDER: '', DATASETS_FOLDER: '', + LORA_INSTALL_PATH: '', }); const [isSettingsLoaded, setIsLoaded] = useState(false); useEffect(() => { @@ -26,6 +28,7 @@ export default function useSettings() { HF_TOKEN: data.HF_TOKEN || '', TRAINING_FOLDER: data.TRAINING_FOLDER || '', DATASETS_FOLDER: data.DATASETS_FOLDER || '', + LORA_INSTALL_PATH: data.LORA_INSTALL_PATH || '', }); setIsLoaded(true); }) diff --git a/ui/src/server/settings.ts b/ui/src/server/settings.ts index 9c6abe68b..a3721b4e8 100644 --- a/ui/src/server/settings.ts +++ b/ui/src/server/settings.ts @@ -67,6 +67,25 @@ export const getHFToken = async () => { return token; }; +export const getLoraInstallPath = async () => { + const key = 'LORA_INSTALL_PATH'; + let loraPath = myCache.get(key) as string; + if (loraPath !== undefined) { + return loraPath; + } + let row = await prisma.settings.findFirst({ + where: { + key: key, + }, + }); + loraPath = ''; + if (row?.value && row.value !== '') { + loraPath = row.value; + } + myCache.set(key, loraPath); + return loraPath; +}; + export const getDataRoot = async () => { const key = 'DATA_ROOT'; let dataRoot = myCache.get(key) as string;