Add ability to run small scripts from the ui and added a merge lora script
This commit is contained in:
Jaret Burkett
277
ui/src/components/MergeLoRAsModal.tsx
Normal file
277
ui/src/components/MergeLoRAsModal.tsx
Normal file
@@ -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;
|
||||
Reference in New Issue
Block a user