278 lines
8.9 KiB
TypeScript
278 lines
8.9 KiB
TypeScript
'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;
|