diff --git a/ui/cron/actions/startJob.ts b/ui/cron/actions/startJob.ts index a6e3983b..8f0eb8bd 100644 --- a/ui/cron/actions/startJob.ts +++ b/ui/cron/actions/startJob.ts @@ -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)) { diff --git a/ui/cron/pythonPath.ts b/ui/cron/pythonPath.ts new file mode 100644 index 00000000..64ac231a --- /dev/null +++ b/ui/cron/pythonPath.ts @@ -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'; +}; diff --git a/ui/src/app/api/scripts/route.ts b/ui/src/app/api/scripts/route.ts new file mode 100644 index 00000000..bab8af49 --- /dev/null +++ b/ui/src/app/api/scripts/route.ts @@ -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)) { + 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 => { + 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 }); +} diff --git a/ui/src/app/layout.tsx b/ui/src/app/layout.tsx index bc692e1b..46747e88 100644 --- a/ui/src/app/layout.tsx +++ b/ui/src/app/layout.tsx @@ -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 }) + ); diff --git a/ui/src/components/FilesWidget.tsx b/ui/src/components/FilesWidget.tsx index 4131a86b..5c5d6632 100644 --- a/ui/src/components/FilesWidget.tsx +++ b/ui/src/components/FilesWidget.tsx @@ -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 }) {

Checkpoints

{files.length} + {files.length > 0 && ( + { + const outputName = `${jobName}_merged`; + openMergeLoRAsModal( + getFoldername(files[0].path), + outputName, + files.map(f => ({ path: f.path })), + () => { + refreshFiles(); + }, + ); + }} + > + merge + + )}
@@ -45,7 +65,7 @@ export default function FilesWidget({ jobID }: { jobID: string }) { {['success', 'refreshing'].includes(status) && (
{files.map((file, index) => { - const fileName = file.path.split('/').pop() || ''; + const fileName = getFilename(file.path); const nameWithoutExt = fileName.replace('.safetensors', ''); return ( {isGPUInfoLoaded && gpuList.length > 0 && }
{jobType === 'train' && (
- +
)}
diff --git a/ui/src/components/MergeLoRAsModal.tsx b/ui/src/components/MergeLoRAsModal.tsx new file mode 100644 index 00000000..ead268ef --- /dev/null +++ b/ui/src/components/MergeLoRAsModal.tsx @@ -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(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([]); + const [isRunning, setIsRunning] = useState(false); + const [isDone, setIsDone] = useState(false); + const [hasError, setHasError] = useState(false); + const [logOutput, setLogOutput] = useState(''); + const logRef = useRef(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 ( + + {showLog ? ( +
+
+ {isRunning && Merging LoRAs... please do not close this window.} + {isDone && hasError && Merge failed. See log below.} + {isDone && !hasError && Merge complete.} +
+
+ {logOutput || (isRunning ? 'Starting...\n' : '')} +
+
+ +
+
+ ) : ( +
{ + e.preventDefault(); + onSubmit(); + }} + > + { + setModalInfo({ + ...modalInfo, + outputName: value, + } as MergeLoRAsModalState); + }} + placeholder="Enter output filename" + /> + +
+ addLoRA(value)} + options={options} + /> +
+ + {selectedLoRAs.length > 0 && ( +
+ +
+ {selectedLoRAs.map(s => ( +
+
+ {loraLabel(s.path)} +
+ { + 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" + /> + +
+ ))} +
+
+ )} + +
+ + +
+ + )} +
+ ); +}; + +export default MergeLoRAsModal; diff --git a/ui/src/components/Modal.tsx b/ui/src/components/Modal.tsx index ed402bd5..ec661f0c 100644 --- a/ui/src/components/Modal.tsx +++ b/ui/src/components/Modal.tsx @@ -68,14 +68,14 @@ export const Modal: React.FC = ({ > {/* Modal panel */}
e.stopPropagation()} > {/* Modal header */} {(title || showCloseButton) && ( -
+
{title && ( - )} diff --git a/ui/src/components/formInputs.tsx b/ui/src/components/formInputs.tsx index c498fa0d..3d8cff6f 100644 --- a/ui/src/components/formInputs.tsx +++ b/ui/src/components/formInputs.tsx @@ -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((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((props: Te )} )} - { - if (!disabled) onChange(e.target.value); - }} - className={`${inputClasses} ${disabled ? 'opacity-30 cursor-not-allowed' : ''}`} - placeholder={placeholder} - required={required} - disabled={disabled} - /> + {suffix ? ( +
+ { + 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} + /> + + {suffix} + +
+ ) : ( + { + if (!disabled) onChange(e.target.value); + }} + className={`${inputClasses} ${disabled ? 'opacity-30 cursor-not-allowed' : ''}`} + placeholder={placeholder} + required={required} + disabled={disabled} + /> + )}
); }); @@ -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 (
{ 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); } }} /> diff --git a/ui/src/utils/basic.ts b/ui/src/utils/basic.ts index 4870a569..09c6ba58 100644 --- a/ui/src/utils/basic.ts +++ b/ui/src/utils/basic.ts @@ -36,6 +36,16 @@ export const objToTags = (obj: Record): 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; diff --git a/ui/src/utils/callScript.ts b/ui/src/utils/callScript.ts new file mode 100644 index 00000000..9ba34290 --- /dev/null +++ b/ui/src/utils/callScript.ts @@ -0,0 +1,127 @@ +import { apiClient } from './api'; + +export type ScriptArgs = (string | number | boolean)[] | Record; + +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 => { + const response = await apiClient.post( + '/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 => { + const token = typeof window !== 'undefined' ? localStorage.getItem('AI_TOOLKIT_AUTH') : null; + const headers: Record = { '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; +}; diff --git a/ui_scripts/merge_loras.py b/ui_scripts/merge_loras.py new file mode 100644 index 00000000..bc844ca0 --- /dev/null +++ b/ui_scripts/merge_loras.py @@ -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()) diff --git a/ui_scripts/test_script.py b/ui_scripts/test_script.py new file mode 100644 index 00000000..ecd8e9b3 --- /dev/null +++ b/ui_scripts/test_script.py @@ -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()) diff --git a/version.py b/version.py index bc5433b5..184e18b0 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.9.10" +VERSION = "0.9.11"