Files
ace-step-ui/server/src/routes/training.ts
T
fspecii 78426c00dd Fix LoRA training Error 500 - remove Radio serialization issue
The @gradio/client was wrapping Radio component values as
{"value": "replace", "__type__": "update"} instead of plain strings,
causing the update_settings Gradio call to fail and crash the server.

- /update-settings: removed Gradio call, settings are applied at save time
- /save-dataset: migrated to REST API (/v1/dataset/save) which accepts
  tag_position and other settings directly as strings
- Frontend: pass dataset settings (tag, position, etc.) when saving
2026-03-02 18:06:44 +02:00

875 lines
30 KiB
TypeScript

import { Router, Request, Response } from 'express';
import { authMiddleware, AuthenticatedRequest } from '../middleware/auth.js';
import { getGradioClient } from '../services/gradio-client.js';
import { config } from '../config/index.js';
import { resolvePythonPath } from '../services/acestep.js';
import multer from 'multer';
import path from 'path';
import { existsSync, readdirSync, statSync, readFileSync } from 'fs';
import { mkdir, writeFile, readFile } from 'fs/promises';
import { execSync, spawn } from 'child_process';
import { randomUUID } from 'crypto';
const router = Router();
// --- Audio upload via multer disk storage ---
const AUDIO_EXTENSIONS = ['.wav', '.mp3', '.flac', '.ogg', '.opus'];
const audioStorage = multer.diskStorage({
destination: async (_req: Request, _file, cb) => {
const datasetName = (_req.body?.datasetName as string) || 'default';
const dest = path.join(config.datasets.uploadsDir, datasetName);
try {
await mkdir(dest, { recursive: true });
cb(null, dest);
} catch (err) {
cb(err as Error, dest);
}
},
filename: (_req, file, cb) => {
// Preserve original filename but ensure uniqueness
const ext = path.extname(file.originalname).toLowerCase();
const base = path.basename(file.originalname, ext);
const safeName = base.replace(/[^a-zA-Z0-9_\-. ]/g, '_');
cb(null, `${safeName}${ext}`);
},
});
const audioUpload = multer({
storage: audioStorage,
limits: { fileSize: 100 * 1024 * 1024 }, // 100MB per file
fileFilter: (_req, file, cb) => {
const ext = path.extname(file.originalname).toLowerCase();
if (AUDIO_EXTENSIONS.includes(ext)) {
cb(null, true);
} else {
cb(new Error(`Unsupported file type: ${ext}. Allowed: ${AUDIO_EXTENSIONS.join(', ')}`));
}
},
});
// Get audio duration via ffprobe
function getAudioDuration(filePath: string): number {
try {
const result = execSync(
`ffprobe -v error -show_entries format=duration -of default=noprint_wrappers=1:nokey=1 "${filePath}"`,
{ encoding: 'utf-8', timeout: 10000 }
);
const duration = parseFloat(result.trim());
return isNaN(duration) ? 0 : Math.round(duration);
} catch {
return 0;
}
}
// Resolve ACE-Step base directory
function getAceStepDir(): string {
const envPath = process.env.ACESTEP_PATH;
if (envPath) {
return path.isAbsolute(envPath) ? envPath : path.resolve(process.cwd(), envPath);
}
return path.resolve(config.datasets.dir, '..');
}
// ================== NEW ROUTES ==================
// POST /api/training/upload-audio — Upload audio files for a dataset
router.post('/upload-audio', authMiddleware, audioUpload.array('audio', 50), async (req: AuthenticatedRequest, res: Response) => {
try {
const files = req.files as Express.Multer.File[];
if (!files || files.length === 0) {
res.status(400).json({ error: 'No audio files uploaded' });
return;
}
const datasetName = (req.body?.datasetName as string) || 'default';
const uploadDir = path.join(config.datasets.uploadsDir, datasetName);
res.json({
files: files.map(f => ({
filename: f.filename,
originalName: f.originalname,
size: f.size,
path: f.path,
})),
uploadDir,
count: files.length,
});
} catch (error) {
console.error('[Training] Upload audio error:', error);
res.status(500).json({ error: error instanceof Error ? error.message : 'Upload failed' });
}
});
// POST /api/training/build-dataset — Scan audio directory + create dataset JSON
router.post('/build-dataset', authMiddleware, async (req: AuthenticatedRequest, res: Response) => {
try {
const {
datasetName = 'my_lora_dataset',
customTag = '',
tagPosition = 'prepend',
allInstrumental = true,
} = req.body;
const audioDir = path.join(config.datasets.uploadsDir, datasetName);
if (!existsSync(audioDir)) {
res.status(400).json({ error: `Audio directory not found: uploads/${datasetName}` });
return;
}
// Scan for audio files
const entries = readdirSync(audioDir);
const audioFiles = entries.filter(f => AUDIO_EXTENSIONS.includes(path.extname(f).toLowerCase()));
if (audioFiles.length === 0) {
res.status(400).json({ error: 'No audio files found in directory' });
return;
}
// Build samples in Gradio's exact format
const samples = audioFiles.map(filename => {
const audioPath = path.join(audioDir, filename);
const duration = getAudioDuration(audioPath);
const baseName = path.basename(filename, path.extname(filename));
// Check for companion .txt lyrics file
let rawLyrics = '';
const lyricsPath = path.join(audioDir, `${baseName}.txt`);
if (existsSync(lyricsPath)) {
try {
rawLyrics = readFileSync(lyricsPath, 'utf-8').trim();
} catch { /* ignore */ }
}
const isInstrumental = allInstrumental || !rawLyrics;
return {
id: randomUUID().slice(0, 8),
audio_path: audioPath,
filename,
caption: '',
genre: '',
lyrics: isInstrumental ? '[Instrumental]' : rawLyrics,
raw_lyrics: rawLyrics,
formatted_lyrics: '',
bpm: null as number | null,
keyscale: '',
timesignature: '',
duration,
language: isInstrumental ? 'instrumental' : 'unknown',
is_instrumental: isInstrumental,
custom_tag: customTag,
labeled: false,
prompt_override: null as string | null,
};
});
// Build dataset JSON
const dataset = {
metadata: {
name: datasetName,
custom_tag: customTag,
tag_position: tagPosition,
created_at: new Date().toISOString(),
num_samples: samples.length,
all_instrumental: allInstrumental,
genre_ratio: 0,
},
samples,
};
// Save JSON to datasets dir
await mkdir(config.datasets.dir, { recursive: true });
const jsonPath = path.join(config.datasets.dir, `${datasetName}.json`);
await writeFile(jsonPath, JSON.stringify(dataset, null, 2), 'utf-8');
// Now load into Gradio state via the existing endpoint
try {
const client = await getGradioClient();
const result = await client.predict('/load_existing_dataset_for_preprocess', [jsonPath]);
const data = result.data as unknown[];
res.json({
status: data[0],
dataframe: data[1],
sampleCount: samples.length,
sample: {
index: data[2],
audio: data[3],
filename: data[4],
caption: data[5],
genre: data[6],
promptOverride: data[7],
lyrics: data[8],
bpm: data[9],
key: data[10],
timeSignature: data[11],
duration: data[12],
language: data[13],
instrumental: data[14],
rawLyrics: data[15],
},
settings: {
datasetName: data[16],
customTag: data[17],
tagPosition: data[18],
allInstrumental: data[19],
genreRatio: data[20],
},
datasetPath: jsonPath,
});
} catch (gradioError) {
// Gradio may not be running — still return dataset info
console.warn('[Training] Gradio load failed, returning dataset JSON only:', gradioError);
res.json({
status: `Dataset saved (${samples.length} samples). Gradio not available for live preview.`,
dataframe: null,
sampleCount: samples.length,
sample: samples.length > 0 ? {
index: 0,
audio: null,
filename: samples[0].filename,
caption: samples[0].caption,
genre: samples[0].genre,
promptOverride: null,
lyrics: samples[0].lyrics,
bpm: samples[0].bpm,
key: samples[0].keyscale,
timeSignature: samples[0].timesignature,
duration: samples[0].duration,
language: samples[0].language,
instrumental: samples[0].is_instrumental,
rawLyrics: samples[0].raw_lyrics,
} : null,
settings: {
datasetName,
customTag,
tagPosition,
allInstrumental,
genreRatio: 0,
},
datasetPath: jsonPath,
});
}
} catch (error) {
console.error('[Training] Build dataset error:', error);
res.status(500).json({ error: error instanceof Error ? error.message : 'Failed to build dataset' });
}
});
// GET /api/training/audio — Proxy audio files from datasets directory
router.get('/audio', authMiddleware, async (req: AuthenticatedRequest, res: Response) => {
try {
let filePath: string;
const aceStepDir = getAceStepDir();
if (req.query.path) {
filePath = req.query.path as string;
} else if (req.query.file) {
// Relative path within datasets dir
filePath = path.join(config.datasets.dir, req.query.file as string);
} else {
res.status(400).json({ error: 'path or file parameter required' });
return;
}
// Path traversal protection
const resolved = path.resolve(filePath);
if (resolved.includes('..') || !resolved.startsWith(aceStepDir)) {
res.status(403).json({ error: 'Access denied: path outside ACE-Step directory' });
return;
}
if (!existsSync(resolved)) {
res.status(404).json({ error: 'Audio file not found' });
return;
}
// Determine content type
const ext = path.extname(resolved).toLowerCase();
const mimeTypes: Record<string, string> = {
'.wav': 'audio/wav',
'.mp3': 'audio/mpeg',
'.flac': 'audio/flac',
'.ogg': 'audio/ogg',
'.opus': 'audio/opus',
};
res.setHeader('Content-Type', mimeTypes[ext] || 'application/octet-stream');
res.sendFile(resolved);
} catch (error) {
console.error('[Training] Audio proxy error:', error);
res.status(500).json({ error: error instanceof Error ? error.message : 'Failed to serve audio' });
}
});
// POST /api/training/preprocess — Spawn Python preprocessing script
router.post('/preprocess', authMiddleware, async (req: AuthenticatedRequest, res: Response) => {
try {
const { datasetPath, outputDir } = req.body;
if (!datasetPath) {
res.status(400).json({ error: 'datasetPath is required' });
return;
}
const aceStepDir = getAceStepDir();
const scriptPath = path.resolve(__dirname, '../../scripts/preprocess_dataset.py');
const pythonPath = resolvePythonPath(aceStepDir);
const resolvedOutput = outputDir || path.join(config.datasets.dir, 'preprocessed_tensors');
// Ensure output dir exists
await mkdir(resolvedOutput, { recursive: true });
// Spawn Python process
const child = spawn(pythonPath, [
scriptPath,
'--dataset', datasetPath,
'--output', resolvedOutput,
'--json',
], {
cwd: aceStepDir,
env: { ...process.env },
});
let stdout = '';
let stderr = '';
child.stdout.on('data', (data: Buffer) => { stdout += data.toString(); });
child.stderr.on('data', (data: Buffer) => { stderr += data.toString(); });
child.on('close', (code: number | null) => {
if (code === 0) {
// Try to parse JSON output
try {
const result = JSON.parse(stdout.trim().split('\n').pop() || '{}');
res.json({ status: 'Preprocessing complete', ...result });
} catch {
res.json({ status: 'Preprocessing complete', output: stdout.trim() });
}
} else {
res.status(500).json({
error: 'Preprocessing failed',
code,
stderr: stderr.trim(),
stdout: stdout.trim(),
});
}
});
child.on('error', (err: Error) => {
res.status(500).json({ error: `Failed to spawn process: ${err.message}` });
});
} catch (error) {
console.error('[Training] Preprocess error:', error);
res.status(500).json({ error: error instanceof Error ? error.message : 'Preprocessing failed' });
}
});
// POST /api/training/scan-directory — Scan a directory for audio files (Node.js implementation)
router.post('/scan-directory', authMiddleware, async (req: AuthenticatedRequest, res: Response) => {
try {
const {
audioDir,
datasetName = 'my_lora_dataset',
customTag = '',
tagPosition = 'prepend',
allInstrumental = true,
} = req.body;
if (!audioDir || typeof audioDir !== 'string') {
res.status(400).json({ error: 'audioDir is required' });
return;
}
// Resolve path — if relative, resolve from ACE-Step dir
const aceStepDir = getAceStepDir();
const resolvedDir = path.isAbsolute(audioDir)
? audioDir
: path.resolve(aceStepDir, audioDir);
if (!existsSync(resolvedDir)) {
res.status(400).json({ error: `Directory not found: ${audioDir}` });
return;
}
// Scan for audio files
const entries = readdirSync(resolvedDir);
const audioFiles = entries.filter(f => AUDIO_EXTENSIONS.includes(path.extname(f).toLowerCase()));
if (audioFiles.length === 0) {
res.status(400).json({ error: 'No audio files found in directory' });
return;
}
// Build table data matching Gradio's format: [#, Filename, Duration, Lyrics, Labeled, BPM, Key, Caption]
const tableHeaders = ['#', 'Filename', 'Duration', 'Lyrics', 'Labeled', 'BPM', 'Key', 'Caption'];
const tableData = audioFiles.map((filename, i) => {
const audioPath = path.join(resolvedDir, filename);
const duration = getAudioDuration(audioPath);
const baseName = path.basename(filename, path.extname(filename));
// Check for companion .txt lyrics file
let lyrics = allInstrumental ? '[Instrumental]' : '';
const lyricsPath = path.join(resolvedDir, `${baseName}.txt`);
if (existsSync(lyricsPath)) {
try {
lyrics = readFileSync(lyricsPath, 'utf-8').trim().slice(0, 50) + '...';
} catch { /* ignore */ }
}
return [i + 1, filename, `${duration}s`, lyrics, '❌', '', '', ''];
});
res.json({
status: `Found ${audioFiles.length} audio files`,
dataframe: {
headers: tableHeaders,
data: tableData,
},
sampleCount: audioFiles.length,
audioDir: resolvedDir,
});
} catch (error) {
console.error('[Training] Scan directory error:', error);
res.status(500).json({ error: error instanceof Error ? error.message : 'Failed to scan directory' });
}
});
// POST /api/training/auto-label — Auto-label dataset samples
// NOTE: Auto-labeling requires the DIT model + LLM to be loaded in Gradio.
// This endpoint attempts to call the Gradio handler. If the Gradio app does not
// expose auto_label_all as a named API, this will fail and the user should use
// the Gradio UI directly.
router.post('/auto-label', authMiddleware, async (req: AuthenticatedRequest, res: Response) => {
try {
const {
skipMetas = false,
formatLyrics = false,
transcribeLyrics = false,
onlyUnlabeled = false,
} = req.body;
// auto_label_all is a lambda-wrapped handler in Gradio, so it may not be accessible
// by name. We try the likely endpoint name; if it fails, return a helpful message.
const client = await getGradioClient();
try {
const result = await client.predict('/auto_label_all', [
skipMetas,
formatLyrics,
transcribeLyrics,
onlyUnlabeled,
]);
const data = result.data as unknown[];
res.json({
dataframe: data[0],
status: data[1],
});
} catch (gradioError) {
// Lambda endpoints aren't named — suggest using Gradio UI
res.status(501).json({
error: 'Auto-labeling requires the Gradio UI. The model must be initialized and the dataset loaded in the Gradio training tab.',
hint: 'Use the Gradio UI at the ACE-Step server URL to auto-label your dataset, then reload it here.',
});
}
} catch (error) {
console.error('[Training] Auto-label error:', error);
res.status(500).json({ error: error instanceof Error ? error.message : 'Auto-label failed' });
}
});
// POST /api/training/init-model — Initialize or change model for training
// NOTE: Model initialization requires the Gradio app. This endpoint attempts to
// call the init_service_wrapper. Since it's a lambda, this may not be accessible.
router.post('/init-model', authMiddleware, async (req: AuthenticatedRequest, res: Response) => {
try {
const {
checkpoint,
configPath,
device = 'auto',
initLlm = false,
lmModelPath = '',
backend = 'pt',
useFlashAttention = false,
offloadToCpu = false,
offloadDitToCpu = false,
compileModel = false,
quantization = false,
} = req.body;
const client = await getGradioClient();
try {
// Try calling by function name (may work if Gradio auto-names it)
const result = await client.predict('/init_service_wrapper', [
checkpoint ?? '',
configPath ?? '',
device,
initLlm,
lmModelPath,
backend,
useFlashAttention,
offloadToCpu,
offloadDitToCpu,
compileModel,
quantization,
]);
const data = result.data as unknown[];
res.json({
status: data[0],
modelReady: !!data[1],
});
} catch (gradioError) {
// Lambda endpoints aren't named — suggest using Gradio UI
res.status(501).json({
error: 'Model initialization requires the Gradio UI.',
hint: 'Initialize the model in the ACE-Step Gradio UI service configuration section, then return here for training.',
});
}
} catch (error) {
console.error('[Training] Init model error:', error);
res.status(500).json({ error: error instanceof Error ? error.message : 'Model init failed' });
}
});
// GET /api/training/checkpoints — List available model checkpoints
router.get('/checkpoints', authMiddleware, async (_req: AuthenticatedRequest, res: Response) => {
try {
const aceStepDir = getAceStepDir();
const checkpointDir = path.join(aceStepDir, 'checkpoints');
if (!existsSync(checkpointDir)) {
res.json({ checkpoints: [], configs: [] });
return;
}
// List checkpoint directories
const entries = readdirSync(checkpointDir);
const checkpoints = entries.filter(e => {
const fullPath = path.join(checkpointDir, e);
return statSync(fullPath).isDirectory();
});
// List config directories (acestep-v15-*)
const configDirs = entries.filter(e =>
e.startsWith('acestep-v15') && statSync(path.join(checkpointDir, e)).isDirectory()
);
res.json({ checkpoints, configs: configDirs });
} catch (error) {
console.error('[Training] List checkpoints error:', error);
res.status(500).json({ error: error instanceof Error ? error.message : 'Failed to list checkpoints' });
}
});
// GET /api/training/lora-checkpoints — List LoRA training checkpoints in output dir
router.get('/lora-checkpoints', authMiddleware, async (req: AuthenticatedRequest, res: Response) => {
try {
const outputDir = (req.query.dir as string) || './lora_output';
const aceStepDir = getAceStepDir();
const resolvedDir = path.isAbsolute(outputDir)
? outputDir
: path.resolve(aceStepDir, outputDir);
if (!existsSync(resolvedDir)) {
res.json({ checkpoints: [] });
return;
}
const entries = readdirSync(resolvedDir);
const checkpointsDir = path.join(resolvedDir, 'checkpoints');
const checkpoints: string[] = [];
if (existsSync(checkpointsDir)) {
const cpEntries = readdirSync(checkpointsDir);
cpEntries.forEach(e => {
if (statSync(path.join(checkpointsDir, e)).isDirectory()) {
checkpoints.push(path.join(checkpointsDir, e));
}
});
}
// Also check for "final" directory
const finalDir = path.join(resolvedDir, 'final');
if (existsSync(finalDir)) {
checkpoints.push(finalDir);
}
res.json({ checkpoints, outputDir: resolvedDir });
} catch (error) {
console.error('[Training] List LoRA checkpoints error:', error);
res.status(500).json({ error: error instanceof Error ? error.message : 'Failed to list checkpoints' });
}
});
// ================== EXISTING ROUTES ==================
// POST /api/training/load-dataset — Load an existing dataset JSON for preprocessing
router.post('/load-dataset', authMiddleware, async (req: AuthenticatedRequest, res: Response) => {
try {
const { datasetPath } = req.body;
if (!datasetPath || typeof datasetPath !== 'string') {
res.status(400).json({ error: 'datasetPath is required' });
return;
}
// Reject path traversal
if (datasetPath.includes('..')) {
res.status(400).json({ error: 'Invalid path' });
return;
}
const client = await getGradioClient();
const result = await client.predict('/load_existing_dataset_for_preprocess', [datasetPath]);
const data = result.data as unknown[];
// Returns: [status, dataframe, sampleIdx, audioPreview, filename, caption, genre,
// promptOverride, lyrics, bpm, key, timesig, duration, language, instrumental,
// rawLyrics, datasetName, customTag, tagPosition, allInstrumental, genreRatio]
res.json({
status: data[0],
dataframe: data[1],
sampleCount: Array.isArray((data[1] as any)?.data) ? (data[1] as any).data.length : 0,
sample: {
index: data[2],
audio: data[3],
filename: data[4],
caption: data[5],
genre: data[6],
promptOverride: data[7],
lyrics: data[8],
bpm: data[9],
key: data[10],
timeSignature: data[11],
duration: data[12],
language: data[13],
instrumental: data[14],
rawLyrics: data[15],
},
settings: {
datasetName: data[16],
customTag: data[17],
tagPosition: data[18],
allInstrumental: data[19],
genreRatio: data[20],
},
});
} catch (error) {
console.error('[Training] Load dataset error:', error);
res.status(500).json({ error: error instanceof Error ? error.message : 'Failed to load dataset' });
}
});
// GET /api/training/sample-preview — Get preview data for a specific sample
router.get('/sample-preview', authMiddleware, async (req: AuthenticatedRequest, res: Response) => {
try {
const idx = parseInt(req.query.idx as string) || 0;
const client = await getGradioClient();
const result = await client.predict('/get_sample_preview', [idx]);
const data = result.data as unknown[];
// Returns: [audio, filename, caption, genre, promptOverride, lyrics, bpm, key, timesig, duration, language, instrumental, rawLyrics]
res.json({
audio: data[0],
filename: data[1],
caption: data[2],
genre: data[3],
promptOverride: data[4],
lyrics: data[5],
bpm: data[6],
key: data[7],
timeSignature: data[8],
duration: data[9],
language: data[10],
instrumental: data[11],
rawLyrics: data[12],
});
} catch (error) {
console.error('[Training] Sample preview error:', error);
res.status(500).json({ error: error instanceof Error ? error.message : 'Failed to get sample preview' });
}
});
// POST /api/training/save-sample — Save edits to a dataset sample
router.post('/save-sample', authMiddleware, async (req: AuthenticatedRequest, res: Response) => {
try {
const { sampleIdx, caption, genre, promptOverride, lyrics, bpm, key, timeSignature, language, instrumental } = req.body;
const client = await getGradioClient();
const result = await client.predict('/save_sample_edit', [
sampleIdx ?? 0,
caption ?? '',
genre ?? '',
promptOverride ?? 'Use Global Ratio',
lyrics ?? '',
bpm ?? 120,
key ?? '',
timeSignature ?? '',
language ?? 'instrumental',
instrumental ?? true,
]);
const data = result.data as unknown[];
// Returns: [dataframe, editStatus]
res.json({
dataframe: data[0],
status: data[1],
});
} catch (error) {
console.error('[Training] Save sample error:', error);
res.status(500).json({ error: error instanceof Error ? error.message : 'Failed to save sample edit' });
}
});
// POST /api/training/update-settings — Update dataset global settings
// Settings are applied directly when saving (via REST API), so no Gradio call needed here.
router.post('/update-settings', authMiddleware, (_req: AuthenticatedRequest, res: Response) => {
res.json({ success: true });
});
// POST /api/training/save-dataset — Save the dataset to a JSON file
router.post('/save-dataset', authMiddleware, async (req: AuthenticatedRequest, res: Response) => {
try {
const { savePath, datasetName, customTag, tagPosition, allInstrumental, genreRatio } = req.body;
const resolvedPath = (savePath ?? `./datasets/${datasetName ?? 'my_lora_dataset'}.json`).trim();
// Use REST API to avoid @gradio/client Radio serialization issues
const apiUrl = config.acestep.apiUrl;
const body: Record<string, unknown> = {
save_path: resolvedPath,
dataset_name: datasetName ?? 'my_lora_dataset',
};
if (customTag !== undefined) body.custom_tag = customTag;
if (tagPosition !== undefined) body.tag_position = tagPosition;
if (allInstrumental !== undefined) body.all_instrumental = allInstrumental;
if (genreRatio !== undefined) body.genre_ratio = genreRatio;
const apiRes = await fetch(`${apiUrl}/v1/dataset/save`, {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify(body),
signal: AbortSignal.timeout(30_000),
});
if (!apiRes.ok) {
const err = await apiRes.json().catch(() => ({})) as any;
throw new Error(err?.detail || err?.error || `Save failed: ${apiRes.status}`);
}
const data = await apiRes.json() as any;
res.json({
status: data.status ?? 'Saved',
path: data.save_path ?? resolvedPath,
});
} catch (error) {
console.error('[Training] Save dataset error:', error);
res.status(500).json({ error: error instanceof Error ? error.message : 'Failed to save dataset' });
}
});
// POST /api/training/load-tensors — Load preprocessed tensors for training
router.post('/load-tensors', authMiddleware, async (req: AuthenticatedRequest, res: Response) => {
try {
const { tensorDir } = req.body;
const client = await getGradioClient();
const result = await client.predict('/load_training_dataset', [
tensorDir ?? './datasets/preprocessed_tensors',
]);
const data = result.data as unknown[];
res.json({ status: data[0] });
} catch (error) {
console.error('[Training] Load tensors error:', error);
res.status(500).json({ error: error instanceof Error ? error.message : 'Failed to load training dataset' });
}
});
// POST /api/training/start — Start LoRA training
router.post('/start', authMiddleware, async (req: AuthenticatedRequest, res: Response) => {
try {
const {
tensorDir, rank, alpha, dropout, learningRate,
epochs, batchSize, gradientAccumulation, saveEvery,
shift, seed, outputDir, resumeCheckpoint,
} = req.body;
const client = await getGradioClient();
const result = await client.predict('/training_wrapper', [
tensorDir ?? './datasets/preprocessed_tensors',
rank ?? 64,
alpha ?? 128,
dropout ?? 0.1,
learningRate ?? 0.0003,
epochs ?? 1000,
batchSize ?? 1,
gradientAccumulation ?? 1,
saveEvery ?? 200,
shift ?? 3.0,
seed ?? 42,
outputDir ?? './lora_output',
resumeCheckpoint ?? null,
]);
const data = result.data as unknown[];
// Returns: [trainingProgress, trainingLog, lineplotData]
res.json({
progress: data[0],
log: data[1],
metrics: data[2],
});
} catch (error) {
console.error('[Training] Start training error:', error);
res.status(500).json({ error: error instanceof Error ? error.message : 'Failed to start training' });
}
});
// POST /api/training/stop — Stop current training
router.post('/stop', authMiddleware, async (_req: AuthenticatedRequest, res: Response) => {
try {
const client = await getGradioClient();
const result = await client.predict('/stop_training', []);
const data = result.data as unknown[];
res.json({ status: data[0] });
} catch (error) {
console.error('[Training] Stop training error:', error);
res.status(500).json({ error: error instanceof Error ? error.message : 'Failed to stop training' });
}
});
// POST /api/training/export — Export trained LoRA weights
router.post('/export', authMiddleware, async (req: AuthenticatedRequest, res: Response) => {
try {
const { exportPath, loraOutputDir } = req.body;
const client = await getGradioClient();
const result = await client.predict('/export_lora', [
exportPath ?? './lora_output/final_lora',
loraOutputDir ?? './lora_output',
]);
const data = result.data as unknown[];
res.json({ status: data[0] });
} catch (error) {
console.error('[Training] Export LoRA error:', error);
res.status(500).json({ error: error instanceof Error ? error.message : 'Failed to export LoRA' });
}
});
// POST /api/training/import-dataset — Import train/test split
router.post('/import-dataset', authMiddleware, async (req: AuthenticatedRequest, res: Response) => {
try {
const { datasetType } = req.body;
const client = await getGradioClient();
const result = await client.predict('/import_dataset', [
datasetType ?? 'train',
]);
const data = result.data as unknown[];
res.json({ status: data[0] });
} catch (error) {
console.error('[Training] Import dataset error:', error);
res.status(500).json({ error: error instanceof Error ? error.message : 'Failed to import dataset' });
}
});
export default router;