Add ability to run small scripts from the ui and added a merge lora script
This commit is contained in:
Jaret Burkett
parent
ec58dcde92
commit
6ecaf679dc
|
|
@ -4,6 +4,7 @@ import { spawn } from 'child_process';
|
|||
import path from 'path';
|
||||
import fs from 'fs';
|
||||
import { TOOLKIT_ROOT, getTrainingFolder, getHFToken } from '../paths';
|
||||
import { resolvePythonPath } from '../pythonPath';
|
||||
const isWindows = process.platform === 'win32';
|
||||
|
||||
const startAndWatchJob = (job: Job) => {
|
||||
|
|
@ -52,21 +53,7 @@ const startAndWatchJob = (job: Job) => {
|
|||
// write the config file
|
||||
fs.writeFileSync(configPath, JSON.stringify(jobConfig, null, 2));
|
||||
|
||||
let pythonPath = 'python';
|
||||
// use .venv or venv if it exists
|
||||
if (fs.existsSync(path.join(TOOLKIT_ROOT, '.venv'))) {
|
||||
if (isWindows) {
|
||||
pythonPath = path.join(TOOLKIT_ROOT, '.venv', 'Scripts', 'python.exe');
|
||||
} else {
|
||||
pythonPath = path.join(TOOLKIT_ROOT, '.venv', 'bin', 'python');
|
||||
}
|
||||
} else if (fs.existsSync(path.join(TOOLKIT_ROOT, 'venv'))) {
|
||||
if (isWindows) {
|
||||
pythonPath = path.join(TOOLKIT_ROOT, 'venv', 'Scripts', 'python.exe');
|
||||
} else {
|
||||
pythonPath = path.join(TOOLKIT_ROOT, 'venv', 'bin', 'python');
|
||||
}
|
||||
}
|
||||
const pythonPath = resolvePythonPath();
|
||||
|
||||
const runFilePath = path.join(TOOLKIT_ROOT, 'run.py');
|
||||
if (!fs.existsSync(runFilePath)) {
|
||||
|
|
|
|||
|
|
@ -0,0 +1,27 @@
|
|||
import path from 'path';
|
||||
import fs from 'fs';
|
||||
import { TOOLKIT_ROOT } from './paths';
|
||||
|
||||
const isWindows = process.platform === 'win32';
|
||||
|
||||
// Shared resolver used by both the cron worker and Next.js API routes
|
||||
// so the Python interpreter is configured in exactly one place.
|
||||
export const resolvePythonPath = (): string => {
|
||||
const candidates: string[] = [];
|
||||
|
||||
if (isWindows) {
|
||||
candidates.push(path.join(TOOLKIT_ROOT, '.venv', 'Scripts', 'python.exe'));
|
||||
candidates.push(path.join(TOOLKIT_ROOT, 'venv', 'Scripts', 'python.exe'));
|
||||
} else {
|
||||
candidates.push(path.join(TOOLKIT_ROOT, '.venv', 'bin', 'python'));
|
||||
candidates.push(path.join(TOOLKIT_ROOT, 'venv', 'bin', 'python'));
|
||||
}
|
||||
|
||||
for (const candidate of candidates) {
|
||||
if (fs.existsSync(candidate)) {
|
||||
return candidate;
|
||||
}
|
||||
}
|
||||
|
||||
return isWindows ? 'python.exe' : 'python3';
|
||||
};
|
||||
|
|
@ -0,0 +1,248 @@
|
|||
import { NextResponse } from 'next/server';
|
||||
import { spawn } from 'child_process';
|
||||
import path from 'path';
|
||||
import fs from 'fs';
|
||||
import { TOOLKIT_ROOT } from '@/paths';
|
||||
import { resolvePythonPath } from '../../../../cron/pythonPath';
|
||||
|
||||
// Long-running scripts: allow up to 20 minutes.
|
||||
export const runtime = 'nodejs';
|
||||
export const dynamic = 'force-dynamic';
|
||||
export const maxDuration = 1200;
|
||||
|
||||
const TIMEOUT_MS = 20 * 60 * 1000;
|
||||
const UI_SCRIPTS_ROOT = path.join(TOOLKIT_ROOT, 'ui_scripts');
|
||||
// Only allow flat script names (no path separators, no traversal).
|
||||
const SCRIPT_NAME_RE = /^[A-Za-z0-9_][A-Za-z0-9_.-]*\.py$/;
|
||||
|
||||
const resolveScriptPath = (rawName: unknown): string | null => {
|
||||
if (typeof rawName !== 'string') return null;
|
||||
const name = rawName.trim();
|
||||
if (!SCRIPT_NAME_RE.test(name)) return null;
|
||||
|
||||
const target = path.resolve(UI_SCRIPTS_ROOT, name);
|
||||
const rootWithSep = UI_SCRIPTS_ROOT.endsWith(path.sep) ? UI_SCRIPTS_ROOT : UI_SCRIPTS_ROOT + path.sep;
|
||||
if (!target.startsWith(rootWithSep)) return null;
|
||||
if (!fs.existsSync(target) || !fs.statSync(target).isFile()) return null;
|
||||
return target;
|
||||
};
|
||||
|
||||
// Args may be a positional list or an object that becomes --key value pairs.
|
||||
// Every value is stringified before being passed to spawn (no shell).
|
||||
const normalizeArgs = (raw: unknown): string[] | { error: string } => {
|
||||
if (raw == null) return [];
|
||||
if (Array.isArray(raw)) {
|
||||
const out: string[] = [];
|
||||
for (const v of raw) {
|
||||
if (v == null) continue;
|
||||
if (typeof v === 'string' || typeof v === 'number' || typeof v === 'boolean') {
|
||||
out.push(String(v));
|
||||
} else {
|
||||
return { error: 'args entries must be string|number|boolean' };
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
if (typeof raw === 'object') {
|
||||
const out: string[] = [];
|
||||
for (const [key, value] of Object.entries(raw as Record<string, unknown>)) {
|
||||
if (!/^[A-Za-z0-9_-]+$/.test(key)) return { error: `invalid arg key: ${key}` };
|
||||
const flag = `--${key}`;
|
||||
if (value === true) {
|
||||
out.push(flag);
|
||||
} else if (value === false || value == null) {
|
||||
continue;
|
||||
} else if (typeof value === 'string' || typeof value === 'number') {
|
||||
out.push(flag, String(value));
|
||||
} else {
|
||||
return { error: `args.${key} must be string|number|boolean` };
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
return { error: 'args must be an array or object' };
|
||||
};
|
||||
|
||||
interface RunResult {
|
||||
ok: boolean;
|
||||
exitCode: number | null;
|
||||
signal: NodeJS.Signals | null;
|
||||
stdout: string;
|
||||
stderr: string;
|
||||
result: unknown;
|
||||
timedOut: boolean;
|
||||
error?: string;
|
||||
}
|
||||
|
||||
// Parses the last line of stdout as JSON if possible — scripts can use this
|
||||
// to return structured data alongside their human-readable logs.
|
||||
const parseResult = (stdout: string): unknown => {
|
||||
const lines = stdout.trimEnd().split(/\r?\n/);
|
||||
for (let i = lines.length - 1; i >= 0; i--) {
|
||||
const line = lines[i].trim();
|
||||
if (!line) continue;
|
||||
if (line.startsWith('{') || line.startsWith('[')) {
|
||||
try {
|
||||
return JSON.parse(line);
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
return null;
|
||||
};
|
||||
|
||||
const runBuffered = (scriptPath: string, args: string[]): Promise<RunResult> => {
|
||||
return new Promise(resolve => {
|
||||
const child = spawn(resolvePythonPath(), ['-u', scriptPath, ...args], {
|
||||
cwd: TOOLKIT_ROOT,
|
||||
env: { ...process.env, PYTHONUNBUFFERED: '1', PYTHONIOENCODING: 'utf-8' },
|
||||
windowsHide: true,
|
||||
});
|
||||
|
||||
let stdout = '';
|
||||
let stderr = '';
|
||||
let timedOut = false;
|
||||
|
||||
const timer = setTimeout(() => {
|
||||
timedOut = true;
|
||||
child.kill('SIGKILL');
|
||||
}, TIMEOUT_MS);
|
||||
|
||||
child.stdout.on('data', (chunk: Buffer) => {
|
||||
stdout += chunk.toString('utf-8');
|
||||
});
|
||||
child.stderr.on('data', (chunk: Buffer) => {
|
||||
stderr += chunk.toString('utf-8');
|
||||
});
|
||||
|
||||
child.on('error', err => {
|
||||
clearTimeout(timer);
|
||||
resolve({
|
||||
ok: false,
|
||||
exitCode: null,
|
||||
signal: null,
|
||||
stdout,
|
||||
stderr,
|
||||
result: null,
|
||||
timedOut,
|
||||
error: err.message,
|
||||
});
|
||||
});
|
||||
|
||||
child.on('close', (code, signal) => {
|
||||
clearTimeout(timer);
|
||||
resolve({
|
||||
ok: !timedOut && code === 0,
|
||||
exitCode: code,
|
||||
signal,
|
||||
stdout,
|
||||
stderr,
|
||||
result: parseResult(stdout),
|
||||
timedOut,
|
||||
error: timedOut ? 'Script timed out after 20 minutes' : undefined,
|
||||
});
|
||||
});
|
||||
});
|
||||
};
|
||||
|
||||
// NDJSON stream: one JSON object per line so clients can parse incrementally.
|
||||
const runStreaming = (scriptPath: string, args: string[]): Response => {
|
||||
const child = spawn(resolvePythonPath(), ['-u', scriptPath, ...args], {
|
||||
cwd: TOOLKIT_ROOT,
|
||||
env: { ...process.env, PYTHONUNBUFFERED: '1', PYTHONIOENCODING: 'utf-8' },
|
||||
windowsHide: true,
|
||||
});
|
||||
|
||||
const encoder = new TextEncoder();
|
||||
let stdoutBuf = '';
|
||||
let stderrBuf = '';
|
||||
let timedOut = false;
|
||||
|
||||
const stream = new ReadableStream({
|
||||
start(controller) {
|
||||
const send = (obj: unknown) => {
|
||||
controller.enqueue(encoder.encode(JSON.stringify(obj) + '\n'));
|
||||
};
|
||||
|
||||
const timer = setTimeout(() => {
|
||||
timedOut = true;
|
||||
send({ type: 'error', message: 'Script timed out after 20 minutes' });
|
||||
child.kill('SIGKILL');
|
||||
}, TIMEOUT_MS);
|
||||
|
||||
child.stdout.on('data', (chunk: Buffer) => {
|
||||
const text = chunk.toString('utf-8');
|
||||
stdoutBuf += text;
|
||||
send({ type: 'stdout', data: text });
|
||||
});
|
||||
child.stderr.on('data', (chunk: Buffer) => {
|
||||
const text = chunk.toString('utf-8');
|
||||
stderrBuf += text;
|
||||
send({ type: 'stderr', data: text });
|
||||
});
|
||||
|
||||
child.on('error', err => {
|
||||
clearTimeout(timer);
|
||||
send({ type: 'error', message: err.message });
|
||||
controller.close();
|
||||
});
|
||||
|
||||
child.on('close', (code, signal) => {
|
||||
clearTimeout(timer);
|
||||
send({
|
||||
type: 'exit',
|
||||
exitCode: code,
|
||||
signal,
|
||||
ok: !timedOut && code === 0,
|
||||
timedOut,
|
||||
result: parseResult(stdoutBuf),
|
||||
stderr: stderrBuf,
|
||||
});
|
||||
controller.close();
|
||||
});
|
||||
},
|
||||
cancel() {
|
||||
if (!child.killed) child.kill('SIGKILL');
|
||||
},
|
||||
});
|
||||
|
||||
return new Response(stream, {
|
||||
headers: {
|
||||
'Content-Type': 'application/x-ndjson; charset=utf-8',
|
||||
'Cache-Control': 'no-cache, no-transform',
|
||||
'X-Accel-Buffering': 'no',
|
||||
},
|
||||
});
|
||||
};
|
||||
|
||||
export async function POST(request: Request) {
|
||||
let body: any;
|
||||
try {
|
||||
body = await request.json();
|
||||
} catch {
|
||||
return NextResponse.json({ error: 'Invalid JSON body' }, { status: 400 });
|
||||
}
|
||||
|
||||
const scriptPath = resolveScriptPath(body?.script);
|
||||
if (!scriptPath) {
|
||||
return NextResponse.json(
|
||||
{ error: 'Invalid or unknown script. Must be a *.py file inside ui_scripts/.' },
|
||||
{ status: 400 },
|
||||
);
|
||||
}
|
||||
|
||||
const normalized = normalizeArgs(body?.args);
|
||||
if (!Array.isArray(normalized)) {
|
||||
return NextResponse.json({ error: normalized.error }, { status: 400 });
|
||||
}
|
||||
|
||||
if (body?.stream === true) {
|
||||
return runStreaming(scriptPath, normalized);
|
||||
}
|
||||
|
||||
const result = await runBuffered(scriptPath, normalized);
|
||||
const status = result.ok ? 200 : result.timedOut ? 504 : 500;
|
||||
return NextResponse.json(result, { status });
|
||||
}
|
||||
|
|
@ -9,6 +9,7 @@ import AuthWrapper from '@/components/AuthWrapper';
|
|||
import DocModal from '@/components/DocModal';
|
||||
import os from 'os';
|
||||
import { CaptionDatasetModal } from '@/components/CaptionDatasetModal';
|
||||
import MergeLoRAsModal from '@/components/MergeLoRAsModal';
|
||||
|
||||
export const dynamic = 'force-dynamic';
|
||||
|
||||
|
|
@ -55,6 +56,7 @@ export default function RootLayout({ children }: { children: React.ReactNode })
|
|||
<ConfirmModal />
|
||||
<DocModal />
|
||||
<CaptionDatasetModal />
|
||||
<MergeLoRAsModal />
|
||||
</body>
|
||||
</html>
|
||||
);
|
||||
|
|
|
|||
|
|
@ -2,8 +2,10 @@ import React from 'react';
|
|||
import useFilesList from '@/hooks/useFilesList';
|
||||
import Link from 'next/link';
|
||||
import { Loader2, AlertCircle, Download, Box, Brain } from 'lucide-react';
|
||||
import { openMergeLoRAsModal } from './MergeLoRAsModal';
|
||||
import { getFilename, getFoldername } from '@/utils/basic';
|
||||
|
||||
export default function FilesWidget({ jobID }: { jobID: string }) {
|
||||
export default function FilesWidget({ jobID, jobName }: { jobID: string; jobName: string }) {
|
||||
const { files, status, refreshFiles } = useFilesList(jobID, 5000);
|
||||
|
||||
const cleanSize = (size: number) => {
|
||||
|
|
@ -26,6 +28,24 @@ export default function FilesWidget({ jobID }: { jobID: string }) {
|
|||
<h2 className="font-semibold text-gray-100">Checkpoints</h2>
|
||||
<span className="px-2 py-0.5 bg-gray-700 rounded-full text-xs text-gray-300">{files.length}</span>
|
||||
</div>
|
||||
{files.length > 0 && (
|
||||
<span
|
||||
className="px-3 py-1 rounded-full text-sm bg-purple-500/10 text-purple-500 uppercase cursor-pointer hover:bg-purple-500/20"
|
||||
onClick={() => {
|
||||
const outputName = `${jobName}_merged`;
|
||||
openMergeLoRAsModal(
|
||||
getFoldername(files[0].path),
|
||||
outputName,
|
||||
files.map(f => ({ path: f.path })),
|
||||
() => {
|
||||
refreshFiles();
|
||||
},
|
||||
);
|
||||
}}
|
||||
>
|
||||
merge
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div className="p-2">
|
||||
|
|
@ -45,7 +65,7 @@ export default function FilesWidget({ jobID }: { jobID: string }) {
|
|||
{['success', 'refreshing'].includes(status) && (
|
||||
<div className="space-y-1">
|
||||
{files.map((file, index) => {
|
||||
const fileName = file.path.split('/').pop() || '';
|
||||
const fileName = getFilename(file.path);
|
||||
const nameWithoutExt = fileName.replace('.safetensors', '');
|
||||
return (
|
||||
<a
|
||||
|
|
|
|||
|
|
@ -170,7 +170,7 @@ export default function JobOverview({ job }: JobOverviewProps) {
|
|||
<div className="mt-4">{isGPUInfoLoaded && gpuList.length > 0 && <GPUWidget gpu={gpuList[0]} />}</div>
|
||||
{jobType === 'train' && (
|
||||
<div className="mt-4">
|
||||
<FilesWidget jobID={job.id} />
|
||||
<FilesWidget jobID={job.id} jobName={job.name} />
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
|
|
|||
|
|
@ -0,0 +1,277 @@
|
|||
'use client';
|
||||
import React, { useEffect, useRef, useState } from 'react';
|
||||
import { createGlobalState } from 'react-global-hooks';
|
||||
import { Trash2 } from 'lucide-react';
|
||||
import { Modal } from './Modal';
|
||||
import { TextInput, SelectInput } from '@/components/formInputs';
|
||||
import { getFilename } from '@/utils/basic';
|
||||
import { callScriptStream } from '@/utils/callScript';
|
||||
import { SelectOption } from '@/types';
|
||||
|
||||
export interface MergeLoRAFile {
|
||||
path: string;
|
||||
}
|
||||
|
||||
export interface MergeLoRAsModalState {
|
||||
folderPath: string;
|
||||
outputName: string;
|
||||
availableLoRAs: MergeLoRAFile[];
|
||||
onClose?: () => void;
|
||||
}
|
||||
|
||||
interface SelectedLoRA {
|
||||
path: string;
|
||||
strength: number;
|
||||
}
|
||||
|
||||
export const mergeLoRAsModalState = createGlobalState<MergeLoRAsModalState | null>(null);
|
||||
|
||||
export const openMergeLoRAsModal = (
|
||||
folderPath: string,
|
||||
outputName: string,
|
||||
availableLoRAs: MergeLoRAFile[],
|
||||
onClose?: () => void,
|
||||
) => {
|
||||
mergeLoRAsModalState.set({
|
||||
folderPath,
|
||||
outputName,
|
||||
availableLoRAs,
|
||||
onClose,
|
||||
});
|
||||
};
|
||||
|
||||
const joinPath = (folder: string, name: string) => {
|
||||
const sep = folder.includes('\\') && !folder.includes('/') ? '\\' : '/';
|
||||
const trimmed = folder.replace(/[\\/]+$/, '');
|
||||
return `${trimmed}${sep}${name}.safetensors`;
|
||||
};
|
||||
|
||||
const MergeLoRAsModal: React.FC = () => {
|
||||
const [modalInfo, setModalInfo] = mergeLoRAsModalState.use();
|
||||
const isOpen = modalInfo !== null;
|
||||
const [selectedLoRAs, setSelectedLoRAs] = useState<SelectedLoRA[]>([]);
|
||||
const [isRunning, setIsRunning] = useState(false);
|
||||
const [isDone, setIsDone] = useState(false);
|
||||
const [hasError, setHasError] = useState(false);
|
||||
const [logOutput, setLogOutput] = useState('');
|
||||
const logRef = useRef<HTMLDivElement | null>(null);
|
||||
|
||||
useEffect(() => {
|
||||
if (!modalInfo) {
|
||||
setSelectedLoRAs([]);
|
||||
setIsRunning(false);
|
||||
setIsDone(false);
|
||||
setHasError(false);
|
||||
setLogOutput('');
|
||||
}
|
||||
}, [modalInfo]);
|
||||
|
||||
useEffect(() => {
|
||||
if (logRef.current) {
|
||||
logRef.current.scrollTop = logRef.current.scrollHeight;
|
||||
}
|
||||
}, [logOutput]);
|
||||
|
||||
const onClose = () => {
|
||||
if (isRunning) return;
|
||||
setModalInfo(null);
|
||||
modalInfo?.onClose?.();
|
||||
};
|
||||
|
||||
const onSubmit = async () => {
|
||||
if (isRunning || !modalInfo) return;
|
||||
setIsRunning(true);
|
||||
setIsDone(false);
|
||||
setHasError(false);
|
||||
setLogOutput('');
|
||||
|
||||
const output = joinPath(modalInfo.folderPath, modalInfo.outputName);
|
||||
|
||||
const append = (chunk: string) => setLogOutput(prev => prev + chunk);
|
||||
|
||||
try {
|
||||
const finalEvent = await callScriptStream('merge_loras.py', {
|
||||
args: {
|
||||
loras: JSON.stringify(selectedLoRAs),
|
||||
output,
|
||||
save_dtype: 'bfloat16',
|
||||
device: 'cpu',
|
||||
},
|
||||
onStdout: append,
|
||||
onStderr: append,
|
||||
});
|
||||
|
||||
const ok = finalEvent?.type === 'exit' && finalEvent.ok === true;
|
||||
if (!ok) {
|
||||
setHasError(true);
|
||||
if (finalEvent?.type === 'error' && finalEvent.message) {
|
||||
append(`\n${finalEvent.message}\n`);
|
||||
} else if (finalEvent?.type === 'exit' && finalEvent.timedOut) {
|
||||
append('\nScript timed out.\n');
|
||||
} else if (finalEvent?.type === 'exit') {
|
||||
append(`\nScript exited with code ${finalEvent.exitCode}.\n`);
|
||||
}
|
||||
}
|
||||
} catch (err: any) {
|
||||
setHasError(true);
|
||||
append(`\n${err?.message || 'Unknown error'}\n`);
|
||||
} finally {
|
||||
setIsRunning(false);
|
||||
setIsDone(true);
|
||||
}
|
||||
};
|
||||
|
||||
const loraLabel = (path: string) => getFilename(path).replace('.safetensors', '');
|
||||
|
||||
const availableLoRAs = modalInfo?.availableLoRAs ?? [];
|
||||
const selectedPaths = selectedLoRAs.map(s => s.path);
|
||||
|
||||
const options: SelectOption[] = availableLoRAs
|
||||
.filter(f => !selectedPaths.includes(f.path))
|
||||
.map(f => ({ value: f.path, label: loraLabel(f.path) }));
|
||||
|
||||
const rescale = (items: SelectedLoRA[]): SelectedLoRA[] => {
|
||||
if (items.length === 0) return items;
|
||||
const strength = Math.round((1 / items.length) * 1000) / 1000;
|
||||
return items.map(s => ({ ...s, strength }));
|
||||
};
|
||||
|
||||
const addLoRA = (path: string) => {
|
||||
if (!path || selectedPaths.includes(path)) return;
|
||||
setSelectedLoRAs(prev => rescale([...prev, { path, strength: 0 }]));
|
||||
};
|
||||
|
||||
const removeLoRA = (path: string) => {
|
||||
setSelectedLoRAs(prev => rescale(prev.filter(s => s.path !== path)));
|
||||
};
|
||||
|
||||
const updateStrength = (path: string, strength: number | null) => {
|
||||
setSelectedLoRAs(prev => prev.map(s => (s.path === path ? { ...s, strength: strength ?? 0 } : s)));
|
||||
};
|
||||
|
||||
const showLog = isRunning || isDone;
|
||||
|
||||
return (
|
||||
<Modal
|
||||
isOpen={isOpen}
|
||||
onClose={onClose}
|
||||
title="Merge LoRAs"
|
||||
size="lg"
|
||||
showCloseButton={!isRunning}
|
||||
closeOnOverlayClick={!isRunning}
|
||||
>
|
||||
{showLog ? (
|
||||
<div>
|
||||
<div className="mb-2 text-sm">
|
||||
{isRunning && <span className="text-amber-400">Merging LoRAs... please do not close this window.</span>}
|
||||
{isDone && hasError && <span className="text-rose-400">Merge failed. See log below.</span>}
|
||||
{isDone && !hasError && <span className="text-emerald-400">Merge complete.</span>}
|
||||
</div>
|
||||
<div
|
||||
ref={logRef}
|
||||
className="font-mono text-xs whitespace-pre-wrap break-all overflow-y-auto rounded-md p-3 min-h-[400px] max-h-[60vh] bg-white text-gray-900 dark:bg-black dark:text-gray-100"
|
||||
>
|
||||
{logOutput || (isRunning ? 'Starting...\n' : '')}
|
||||
</div>
|
||||
<div className="mt-4 flex justify-end gap-2">
|
||||
<button
|
||||
type="button"
|
||||
onClick={onClose}
|
||||
disabled={isRunning}
|
||||
className="px-4 py-2 text-sm bg-gray-700 hover:bg-gray-600 disabled:opacity-40 disabled:cursor-not-allowed text-gray-100 rounded-md"
|
||||
>
|
||||
Close
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
) : (
|
||||
<form
|
||||
onSubmit={e => {
|
||||
e.preventDefault();
|
||||
onSubmit();
|
||||
}}
|
||||
>
|
||||
<TextInput
|
||||
label="Output Filename"
|
||||
value={modalInfo?.outputName || ''}
|
||||
suffix=".safetensors"
|
||||
onChange={value => {
|
||||
setModalInfo({
|
||||
...modalInfo,
|
||||
outputName: value,
|
||||
} as MergeLoRAsModalState);
|
||||
}}
|
||||
placeholder="Enter output filename"
|
||||
/>
|
||||
|
||||
<div className="mt-4">
|
||||
<SelectInput
|
||||
label="Add LoRA"
|
||||
multiple={false}
|
||||
value=""
|
||||
onChange={value => addLoRA(value)}
|
||||
options={options}
|
||||
/>
|
||||
</div>
|
||||
|
||||
{selectedLoRAs.length > 0 && (
|
||||
<div className="mt-4">
|
||||
<label className="block text-xs mb-1 text-gray-300">Selected LoRAs</label>
|
||||
<div className="bg-purple-500/10 rounded-xl p-2 max-h-48 overflow-y-auto space-y-1">
|
||||
{selectedLoRAs.map(s => (
|
||||
<div key={s.path} className="flex items-center gap-2 px-2 py-0.5">
|
||||
<div
|
||||
className="flex-1 min-w-0 text-xs text-gray-200 overflow-hidden text-ellipsis whitespace-nowrap"
|
||||
title={loraLabel(s.path)}
|
||||
>
|
||||
{loraLabel(s.path)}
|
||||
</div>
|
||||
<input
|
||||
type="number"
|
||||
value={s.strength}
|
||||
onChange={e => {
|
||||
const raw = e.target.value;
|
||||
if (raw === '' || raw === '-') return;
|
||||
const n = Number(raw);
|
||||
if (!isNaN(n)) updateStrength(s.path, n);
|
||||
}}
|
||||
step="any"
|
||||
className="w-20 flex-shrink-0 text-xs px-2 py-0.5 bg-gray-950 dark:bg-gray-800 border border-gray-700 rounded-sm text-gray-100 focus:ring-1 focus:ring-gray-600 focus:outline-none"
|
||||
/>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => removeLoRA(s.path)}
|
||||
className="flex-shrink-0 text-gray-400 hover:text-rose-400 p-0.5"
|
||||
aria-label="Remove"
|
||||
>
|
||||
<Trash2 className="w-3.5 h-3.5" />
|
||||
</button>
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
<div className="mt-6 flex justify-end gap-2">
|
||||
<button
|
||||
type="button"
|
||||
onClick={onClose}
|
||||
className="px-4 py-2 text-sm text-gray-300 hover:text-gray-100 rounded-md"
|
||||
>
|
||||
Cancel
|
||||
</button>
|
||||
<button
|
||||
type="submit"
|
||||
disabled={selectedLoRAs.length === 0 || !modalInfo?.outputName}
|
||||
className="px-4 py-2 text-sm bg-purple-600 hover:bg-purple-700 disabled:opacity-40 disabled:cursor-not-allowed text-white rounded-md"
|
||||
>
|
||||
Merge
|
||||
</button>
|
||||
</div>
|
||||
</form>
|
||||
)}
|
||||
</Modal>
|
||||
);
|
||||
};
|
||||
|
||||
export default MergeLoRAsModal;
|
||||
|
|
@ -68,14 +68,14 @@ export const Modal: React.FC<ModalProps> = ({
|
|||
>
|
||||
{/* Modal panel */}
|
||||
<div
|
||||
className={`relative mx-auto w-full ${sizeClasses[size]} rounded-lg bg-white dark:bg-gray-900 border border-gray-700 shadow-xl transition-all`}
|
||||
className={`relative mx-auto w-full ${sizeClasses[size]} rounded-xl bg-gray-900 shadow-xl shadow-black/10 dark:shadow-2xl dark:shadow-black/80 transition-all`}
|
||||
onClick={e => e.stopPropagation()}
|
||||
>
|
||||
{/* Modal header */}
|
||||
{(title || showCloseButton) && (
|
||||
<div className="flex items-center justify-between rounded-t-lg border-b border-gray-700 bg-gray-850 px-6 py-4">
|
||||
<div className="flex items-center justify-between bg-gray-800 px-4 py-3 rounded-t-xl">
|
||||
{title && (
|
||||
<h3 id="modal-title" className="text-xl font-semibold text-gray-100">
|
||||
<h3 id="modal-title" className="font-semibold text-gray-100">
|
||||
{title}
|
||||
</h3>
|
||||
)}
|
||||
|
|
|
|||
|
|
@ -28,10 +28,22 @@ export interface TextInputProps extends InputProps {
|
|||
onChange: (value: string) => void;
|
||||
type?: 'text' | 'password';
|
||||
disabled?: boolean;
|
||||
suffix?: React.ReactNode;
|
||||
}
|
||||
|
||||
export const TextInput = forwardRef<HTMLInputElement, TextInputProps>((props: TextInputProps, ref) => {
|
||||
const { label, value, onChange, placeholder, required, disabled, type = 'text', className, docKey = null } = props;
|
||||
const {
|
||||
label,
|
||||
value,
|
||||
onChange,
|
||||
placeholder,
|
||||
required,
|
||||
disabled,
|
||||
type = 'text',
|
||||
className,
|
||||
docKey = null,
|
||||
suffix,
|
||||
} = props;
|
||||
let { doc } = props;
|
||||
if (!doc && docKey) {
|
||||
doc = getDoc(docKey);
|
||||
|
|
@ -48,18 +60,43 @@ export const TextInput = forwardRef<HTMLInputElement, TextInputProps>((props: Te
|
|||
)}
|
||||
</label>
|
||||
)}
|
||||
<input
|
||||
ref={ref}
|
||||
type={type}
|
||||
value={value}
|
||||
onChange={e => {
|
||||
if (!disabled) onChange(e.target.value);
|
||||
}}
|
||||
className={`${inputClasses} ${disabled ? 'opacity-30 cursor-not-allowed' : ''}`}
|
||||
placeholder={placeholder}
|
||||
required={required}
|
||||
disabled={disabled}
|
||||
/>
|
||||
{suffix ? (
|
||||
<div
|
||||
className={classNames(
|
||||
'flex items-stretch w-full bg-gray-950 dark:bg-gray-800 border border-gray-700 rounded-sm focus-within:ring-2 focus-within:ring-gray-600',
|
||||
disabled ? 'opacity-30 cursor-not-allowed' : '',
|
||||
)}
|
||||
>
|
||||
<input
|
||||
ref={ref}
|
||||
type={type}
|
||||
value={value}
|
||||
onChange={e => {
|
||||
if (!disabled) onChange(e.target.value);
|
||||
}}
|
||||
className="flex-1 min-w-0 bg-transparent text-sm px-3 py-1 text-gray-100 placeholder:text-gray-500 focus:outline-none"
|
||||
placeholder={placeholder}
|
||||
required={required}
|
||||
disabled={disabled}
|
||||
/>
|
||||
<span className="flex items-center px-2 text-sm text-gray-400 border-l border-gray-700 bg-gray-900/50 select-none">
|
||||
{suffix}
|
||||
</span>
|
||||
</div>
|
||||
) : (
|
||||
<input
|
||||
ref={ref}
|
||||
type={type}
|
||||
value={value}
|
||||
onChange={e => {
|
||||
if (!disabled) onChange(e.target.value);
|
||||
}}
|
||||
className={`${inputClasses} ${disabled ? 'opacity-30 cursor-not-allowed' : ''}`}
|
||||
placeholder={placeholder}
|
||||
required={required}
|
||||
disabled={disabled}
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
});
|
||||
|
|
@ -187,30 +224,43 @@ export const NumberInput = (props: NumberInputProps) => {
|
|||
);
|
||||
};
|
||||
|
||||
export interface SelectInputProps extends InputProps {
|
||||
value: string;
|
||||
interface SelectInputPropsBase extends InputProps {
|
||||
disabled?: boolean;
|
||||
onChange: (value: string) => void;
|
||||
options: GroupedSelectOption[] | SelectOption[];
|
||||
}
|
||||
|
||||
export interface SingleSelectInputProps extends SelectInputPropsBase {
|
||||
multiple?: false;
|
||||
value: string;
|
||||
onChange: (value: string) => void;
|
||||
}
|
||||
|
||||
export interface MultiSelectInputProps extends SelectInputPropsBase {
|
||||
multiple: true;
|
||||
value: string[];
|
||||
onChange: (value: string[]) => void;
|
||||
}
|
||||
|
||||
export type SelectInputProps = SingleSelectInputProps | MultiSelectInputProps;
|
||||
|
||||
export const SelectInput = (props: SelectInputProps) => {
|
||||
const { label, value, onChange, options, docKey = null } = props;
|
||||
const { label, value, onChange, options, docKey = null, multiple } = props;
|
||||
let { doc } = props;
|
||||
if (!doc && docKey) {
|
||||
doc = getDoc(docKey);
|
||||
}
|
||||
let selectedOption: SelectOption | undefined;
|
||||
if (options && options.length > 0) {
|
||||
// see if grouped options
|
||||
if ('options' in options[0]) {
|
||||
selectedOption = (options as GroupedSelectOption[])
|
||||
.flatMap(group => group.options)
|
||||
.find(opt => opt.value === value);
|
||||
} else {
|
||||
selectedOption = (options as SelectOption[]).find(opt => opt.value === value);
|
||||
}
|
||||
}
|
||||
|
||||
const flatOptions: SelectOption[] =
|
||||
options && options.length > 0
|
||||
? 'options' in options[0]
|
||||
? (options as GroupedSelectOption[]).flatMap(group => group.options)
|
||||
: (options as SelectOption[])
|
||||
: [];
|
||||
|
||||
const selectedOption = multiple
|
||||
? flatOptions.filter(opt => (value as string[]).includes(opt.value))
|
||||
: flatOptions.find(opt => opt.value === (value as string));
|
||||
|
||||
return (
|
||||
<div
|
||||
className={classNames(props.className, {
|
||||
|
|
@ -231,11 +281,15 @@ export const SelectInput = (props: SelectInputProps) => {
|
|||
value={selectedOption}
|
||||
options={options}
|
||||
isDisabled={props.disabled}
|
||||
isMulti={multiple}
|
||||
className="aitk-react-select-container"
|
||||
classNamePrefix="aitk-react-select"
|
||||
onChange={selected => {
|
||||
if (selected) {
|
||||
onChange((selected as { value: string }).value);
|
||||
if (multiple) {
|
||||
const arr = (selected as { value: string }[] | null) ?? [];
|
||||
(onChange as (v: string[]) => void)(arr.map(o => o.value));
|
||||
} else if (selected) {
|
||||
(onChange as (v: string) => void)((selected as { value: string }).value);
|
||||
}
|
||||
}}
|
||||
/>
|
||||
|
|
|
|||
|
|
@ -36,6 +36,16 @@ export const objToTags = (obj: Record<string, any>): string => {
|
|||
.join('\n');
|
||||
};
|
||||
|
||||
export const getFilename = (filePath: string) => {
|
||||
const idx = Math.max(filePath.lastIndexOf('/'), filePath.lastIndexOf('\\'));
|
||||
return idx === -1 ? filePath : filePath.slice(idx + 1);
|
||||
};
|
||||
|
||||
export const getFoldername = (filePath: string) => {
|
||||
const idx = Math.max(filePath.lastIndexOf('/'), filePath.lastIndexOf('\\'));
|
||||
return idx === -1 ? '' : filePath.slice(0, idx);
|
||||
};
|
||||
|
||||
export const pathJoin = (...parts: string[]) => {
|
||||
const sep = parts.length > 0 && parts[0].includes('\\') ? '\\' : '/';
|
||||
const leadingTrailing = sep === '\\' ? /^\\+|\\+$/g : /^\/+|\/+$/g;
|
||||
|
|
|
|||
|
|
@ -0,0 +1,127 @@
|
|||
import { apiClient } from './api';
|
||||
|
||||
export type ScriptArgs = (string | number | boolean)[] | Record<string, string | number | boolean | null | undefined>;
|
||||
|
||||
export interface ScriptResult {
|
||||
ok: boolean;
|
||||
exitCode: number | null;
|
||||
signal: string | null;
|
||||
stdout: string;
|
||||
stderr: string;
|
||||
result: unknown;
|
||||
timedOut: boolean;
|
||||
error?: string;
|
||||
}
|
||||
|
||||
export interface StreamEvent {
|
||||
type: 'stdout' | 'stderr' | 'exit' | 'error';
|
||||
data?: string;
|
||||
message?: string;
|
||||
exitCode?: number | null;
|
||||
signal?: string | null;
|
||||
ok?: boolean;
|
||||
timedOut?: boolean;
|
||||
result?: unknown;
|
||||
stderr?: string;
|
||||
}
|
||||
|
||||
export interface CallScriptOptions {
|
||||
args?: ScriptArgs;
|
||||
signal?: AbortSignal;
|
||||
// Match the API's 20-minute ceiling by default so axios doesn't bail early.
|
||||
timeoutMs?: number;
|
||||
}
|
||||
|
||||
export interface StreamCallScriptOptions extends CallScriptOptions {
|
||||
onEvent?: (event: StreamEvent) => void;
|
||||
onStdout?: (chunk: string) => void;
|
||||
onStderr?: (chunk: string) => void;
|
||||
}
|
||||
|
||||
const DEFAULT_TIMEOUT_MS = 20 * 60 * 1000;
|
||||
|
||||
// Buffered call: resolves with full stdout/stderr after the script exits.
|
||||
export const callScript = async (
|
||||
script: string,
|
||||
options: CallScriptOptions = {},
|
||||
): Promise<ScriptResult> => {
|
||||
const response = await apiClient.post<ScriptResult>(
|
||||
'/api/scripts',
|
||||
{ script, args: options.args },
|
||||
{
|
||||
timeout: options.timeoutMs ?? DEFAULT_TIMEOUT_MS,
|
||||
signal: options.signal,
|
||||
// Don't throw on non-2xx so callers can inspect the structured failure.
|
||||
validateStatus: () => true,
|
||||
},
|
||||
);
|
||||
return response.data;
|
||||
};
|
||||
|
||||
// Streaming call: invokes the callbacks for each NDJSON event as it arrives,
|
||||
// then resolves with the final exit event.
|
||||
export const callScriptStream = async (
|
||||
script: string,
|
||||
options: StreamCallScriptOptions = {},
|
||||
): Promise<StreamEvent | null> => {
|
||||
const token = typeof window !== 'undefined' ? localStorage.getItem('AI_TOOLKIT_AUTH') : null;
|
||||
const headers: Record<string, string> = { 'Content-Type': 'application/json' };
|
||||
if (token) headers['Authorization'] = `Bearer ${token}`;
|
||||
|
||||
const controller = new AbortController();
|
||||
const onAbort = () => controller.abort();
|
||||
if (options.signal) {
|
||||
if (options.signal.aborted) controller.abort();
|
||||
else options.signal.addEventListener('abort', onAbort);
|
||||
}
|
||||
|
||||
const timeout = options.timeoutMs ?? DEFAULT_TIMEOUT_MS;
|
||||
const timer = setTimeout(() => controller.abort(), timeout);
|
||||
|
||||
let finalEvent: StreamEvent | null = null;
|
||||
|
||||
try {
|
||||
const response = await fetch('/api/scripts', {
|
||||
method: 'POST',
|
||||
headers,
|
||||
body: JSON.stringify({ script, args: options.args, stream: true }),
|
||||
signal: controller.signal,
|
||||
});
|
||||
|
||||
if (!response.ok || !response.body) {
|
||||
const text = await response.text().catch(() => '');
|
||||
throw new Error(`Script stream failed: HTTP ${response.status} ${text}`);
|
||||
}
|
||||
|
||||
const reader = response.body.getReader();
|
||||
const decoder = new TextDecoder('utf-8');
|
||||
let buffer = '';
|
||||
|
||||
while (true) {
|
||||
const { value, done } = await reader.read();
|
||||
if (done) break;
|
||||
buffer += decoder.decode(value, { stream: true });
|
||||
|
||||
let newlineIdx: number;
|
||||
while ((newlineIdx = buffer.indexOf('\n')) >= 0) {
|
||||
const line = buffer.slice(0, newlineIdx).trim();
|
||||
buffer = buffer.slice(newlineIdx + 1);
|
||||
if (!line) continue;
|
||||
try {
|
||||
const event = JSON.parse(line) as StreamEvent;
|
||||
options.onEvent?.(event);
|
||||
if (event.type === 'stdout' && event.data) options.onStdout?.(event.data);
|
||||
if (event.type === 'stderr' && event.data) options.onStderr?.(event.data);
|
||||
if (event.type === 'exit' || event.type === 'error') finalEvent = event;
|
||||
} catch {
|
||||
// Ignore malformed lines and keep streaming.
|
||||
}
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
clearTimeout(timer);
|
||||
if (options.signal) options.signal.removeEventListener('abort', onAbort);
|
||||
}
|
||||
|
||||
return finalEvent;
|
||||
};
|
||||
|
|
@ -0,0 +1,134 @@
|
|||
"""Merge a list of LoRAs into a single checkpoint."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
import torch
|
||||
from safetensors.torch import load_file, save_file
|
||||
from safetensors import safe_open
|
||||
|
||||
|
||||
DTYPE_MAP = {
|
||||
"float32": torch.float32,
|
||||
"fp32": torch.float32,
|
||||
"float16": torch.float16,
|
||||
"fp16": torch.float16,
|
||||
"bfloat16": torch.bfloat16,
|
||||
"bf16": torch.bfloat16,
|
||||
}
|
||||
|
||||
|
||||
def log(message: str) -> None:
|
||||
print(message, flush=True)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Merge a list of LoRAs into a single checkpoint."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--loras",
|
||||
required=True,
|
||||
help='JSON list of {"path": "...", "strength": 1.0} entries.',
|
||||
)
|
||||
parser.add_argument("--output", required=True, help="Output .safetensors path.")
|
||||
parser.add_argument(
|
||||
"--save_dtype",
|
||||
default="bfloat16",
|
||||
choices=list(DTYPE_MAP.keys()),
|
||||
help="Dtype of the saved tensors (merging is always done in float32).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--device",
|
||||
default="cpu",
|
||||
help="Device to merge on (cpu, cuda, cuda:1, mps). Output is always saved from CPU.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
try:
|
||||
loras = json.loads(args.loras)
|
||||
except json.JSONDecodeError as e:
|
||||
print(f"Failed to parse --loras JSON: {e}", file=sys.stderr, flush=True)
|
||||
return 2
|
||||
|
||||
if not isinstance(loras, list) or len(loras) == 0:
|
||||
print("--loras must be a non-empty JSON list.", file=sys.stderr, flush=True)
|
||||
return 2
|
||||
|
||||
device = torch.device(args.device)
|
||||
save_dtype = DTYPE_MAP[args.save_dtype]
|
||||
|
||||
log(f"Merging {len(loras)} LoRA(s) on {device}, saving as {args.save_dtype}.")
|
||||
|
||||
merged: dict[str, torch.Tensor] = {}
|
||||
|
||||
metadata = {}
|
||||
|
||||
for i, entry in enumerate(loras):
|
||||
if not isinstance(entry, dict) or "path" not in entry:
|
||||
print(
|
||||
f"LoRA entry {i} must be an object with a 'path' field.",
|
||||
file=sys.stderr,
|
||||
flush=True,
|
||||
)
|
||||
return 2
|
||||
|
||||
path = entry["path"]
|
||||
strength = float(entry.get("strength", 1.0))
|
||||
|
||||
if not os.path.isfile(path):
|
||||
print(f"LoRA file not found: {path}", file=sys.stderr, flush=True)
|
||||
return 2
|
||||
|
||||
log(f"[{i + 1}/{len(loras)}] Loading {path} (strength={strength})")
|
||||
state_dict = load_file(path, device=str(device))
|
||||
|
||||
for key, tensor in state_dict.items():
|
||||
scaled = tensor.to(torch.float32) * strength
|
||||
if key in merged:
|
||||
merged[key].add_(scaled)
|
||||
else:
|
||||
merged[key] = scaled
|
||||
del state_dict
|
||||
|
||||
if i == 0:
|
||||
# For the first LoRA, also copy over all non-tensor metadata (e.g. base model info)
|
||||
with safe_open(path, framework="pt") as f:
|
||||
metadata_to_keep = [
|
||||
"version",
|
||||
"format",
|
||||
"ss_base_model_version",
|
||||
"software",
|
||||
]
|
||||
orig_metadata = f.metadata()
|
||||
for meta_key in metadata_to_keep:
|
||||
if meta_key in orig_metadata:
|
||||
metadata[meta_key] = orig_metadata[meta_key]
|
||||
|
||||
log(f"Casting to {args.save_dtype} and moving to CPU")
|
||||
final = {k: v.to(save_dtype).cpu().contiguous() for k, v in merged.items()}
|
||||
merged.clear()
|
||||
|
||||
log(f"Saving merged checkpoint to {args.output}")
|
||||
save_file(final, args.output, metadata=metadata)
|
||||
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"ok": True,
|
||||
"output": args.output,
|
||||
"num_loras": len(loras),
|
||||
"num_keys": len(final),
|
||||
"save_dtype": args.save_dtype,
|
||||
"device": str(device),
|
||||
}
|
||||
),
|
||||
flush=True,
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
|
|
@ -0,0 +1,46 @@
|
|||
"""Example UI script invoked via POST /api/scripts.
|
||||
|
||||
The API streams stdout/stderr back to the caller. To return structured data,
|
||||
print a single JSON object on the final line of stdout — the route parses it
|
||||
into the `result` field of the response.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description="Example ui_scripts entry point.")
|
||||
parser.add_argument("--message", default="hello", help="Message to echo back.")
|
||||
parser.add_argument(
|
||||
"--count", type=int, default=3, help="Number of log lines to emit."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--delay", type=float, default=0.0, help="Seconds to sleep between log lines."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--fail", action="store_true", help="Exit non-zero to demo failure handling."
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
for i in range(args.count):
|
||||
print(f"[{i + 1}/{args.count}] {args.message}", flush=True)
|
||||
if args.delay > 0:
|
||||
time.sleep(args.delay)
|
||||
|
||||
if args.fail:
|
||||
print("intentional failure", file=sys.stderr, flush=True)
|
||||
return 1
|
||||
|
||||
# Final line: JSON payload the API surfaces as `result`.
|
||||
print(
|
||||
json.dumps({"message": args.message, "count": args.count, "ok": True}),
|
||||
flush=True,
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
|
|
@ -1 +1 @@
|
|||
VERSION = "0.9.10"
|
||||
VERSION = "0.9.11"
|
||||
|
|
|
|||
Loading…
Reference in New Issue