Add ability to run small scripts from the ui and added a merge lora script

2026-05-18 14:33:47 -06:00
parent ec58dcde92
commit 6ecaf679dc
14 changed files with 984 additions and 52 deletions
--- 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 }) {
          <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
--- a/ui/src/components/JobOverview.tsx
+++ b/ui/src/components/JobOverview.tsx
@@ -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>
--- a/ui/src/components/MergeLoRAsModal.tsx
+++ 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<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;
--- a/ui/src/components/Modal.tsx
+++ b/ui/src/components/Modal.tsx
@@ -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>
              )}
--- 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<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);
          }
        }}
      />