From 43989cc19e6e0a32bddfe45e18f96b67f1d86ee2 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Tue, 28 Apr 2026 10:49:46 -0600 Subject: [PATCH] Add advanced config section to the captioner --- ui/src/app/jobs/new/page.tsx | 27 ++- .../AdvancedConfigEditor.tsx} | 62 ++---- ui/src/components/CaptionDatasetModal.tsx | 210 ++++-------------- ui/src/components/CaptionSimpleJob.tsx | 184 +++++++++++++++ 4 files changed, 264 insertions(+), 219 deletions(-) rename ui/src/{app/jobs/new/AdvancedJob.tsx => components/AdvancedConfigEditor.tsx} (69%) create mode 100644 ui/src/components/CaptionSimpleJob.tsx diff --git a/ui/src/app/jobs/new/page.tsx b/ui/src/app/jobs/new/page.tsx index c106a4b0..71a616d5 100644 --- a/ui/src/app/jobs/new/page.tsx +++ b/ui/src/app/jobs/new/page.tsx @@ -17,7 +17,7 @@ import { TopBar, MainContent } from '@/components/layout'; import { Button } from '@headlessui/react'; import { FaChevronLeft } from 'react-icons/fa'; import SimpleJob from './SimpleJob'; -import AdvancedJob from './AdvancedJob'; +import AdvancedConfigEditor from '@/components/AdvancedConfigEditor'; import ErrorBoundary from '@/components/ErrorBoundary'; import { apiClient } from '@/utils/api'; @@ -279,17 +279,20 @@ export default function TrainingForm() { {showAdvancedView ? (
- { + try { + parsed.config.process[0].sqlite_db_path = './aitk_db.db'; + parsed.config.process[0].training_folder = settings.TRAINING_FOLDER; + parsed.config.process[0].device = 'cuda'; + parsed.config.process[0].performance_log_every = 10; + } catch (e) { + console.warn(e); + } + return migrateJobConfig(parsed); + }} />
) : ( diff --git a/ui/src/app/jobs/new/AdvancedJob.tsx b/ui/src/components/AdvancedConfigEditor.tsx similarity index 69% rename from ui/src/app/jobs/new/AdvancedJob.tsx rename to ui/src/components/AdvancedConfigEditor.tsx index b893e62d..ad6f916b 100644 --- a/ui/src/app/jobs/new/AdvancedJob.tsx +++ b/ui/src/components/AdvancedConfigEditor.tsx @@ -1,28 +1,16 @@ 'use client'; import { useEffect, useState, useRef } from 'react'; -import { JobConfig } from '@/types'; import YAML from 'yaml'; import Editor, { OnMount } from '@monaco-editor/react'; import type { editor } from 'monaco-editor'; -import { Settings } from '@/hooks/useSettings'; -import { migrateJobConfig } from './jobConfig'; import { useTheme } from '@/components/ThemeProvider'; -type Props = { - jobConfig: JobConfig; - setJobConfig: (value: any, key?: string) => void; - status: 'idle' | 'saving' | 'success' | 'error'; - handleSubmit: (event: React.FormEvent) => void; - runId: string | null; - gpuIDs: string | null; - setGpuIDs: (value: string | null) => void; - gpuList: any; - datasetOptions: any; - settings: Settings; +type Props = { + config: T; + setConfig: (value: any, key?: string) => void; + transformOnParse?: (parsed: any) => any; }; -const isDev = process.env.NODE_ENV === 'development'; - const yamlConfig: YAML.DocumentOptions & YAML.SchemaOptions & YAML.ParseOptions & @@ -47,11 +35,11 @@ function toYaml(obj: any): string { return doc.toString(yamlConfig); } -export default function AdvancedJob({ jobConfig, setJobConfig, settings }: Props) { +export default function AdvancedConfigEditor({ config, setConfig, transformOnParse }: Props) { const { theme } = useTheme(); const [editorValue, setEditorValue] = useState(''); const [hasError, setHasError] = useState(false); - const lastJobConfigUpdateStringRef = useRef(''); + const lastConfigUpdateStringRef = useRef(''); const editorRef = useRef(null); const monacoRef = useRef(null); @@ -66,17 +54,17 @@ export default function AdvancedJob({ jobConfig, setJobConfig, settings }: Props // Initial content setup try { - const yamlContent = toYaml(jobConfig); + const yamlContent = toYaml(config); setEditorValue(yamlContent); - lastJobConfigUpdateStringRef.current = JSON.stringify(jobConfig); + lastConfigUpdateStringRef.current = JSON.stringify(config); } catch (e) { console.warn(e); } }; useEffect(() => { - const lastUpdate = lastJobConfigUpdateStringRef.current; - const currentUpdate = JSON.stringify(jobConfig); + const lastUpdate = lastConfigUpdateStringRef.current; + const currentUpdate = JSON.stringify(config); // Skip if no changes or editor not yet mounted if (lastUpdate === currentUpdate || !isEditorMounted.current) { @@ -93,7 +81,7 @@ export default function AdvancedJob({ jobConfig, setJobConfig, settings }: Props const scrollTop = editor.getScrollTop(); // Update content - const yamlContent = toYaml(jobConfig); + const yamlContent = toYaml(config); // Only update if the content is actually different if (yamlContent !== editor.getValue()) { @@ -106,12 +94,12 @@ export default function AdvancedJob({ jobConfig, setJobConfig, settings }: Props editor.setScrollTop(scrollTop); } - lastJobConfigUpdateStringRef.current = currentUpdate; + lastConfigUpdateStringRef.current = currentUpdate; } } catch (e) { console.warn(e); } - }, [jobConfig]); + }, [config]); const setMarkers = (errors: { message: string; line: number }[]) => { const monaco = monacoRef.current; @@ -132,27 +120,18 @@ export default function AdvancedJob({ jobConfig, setJobConfig, settings }: Props if (value === undefined) return; try { - const parsed = YAML.parse(value); + let parsed = YAML.parse(value); setHasError(false); setMarkers([]); - // Don't update jobConfig if the change came from the editor itself + // Don't update config if the change came from the editor itself // to avoid a circular update loop - if (JSON.stringify(parsed) !== lastJobConfigUpdateStringRef.current) { - lastJobConfigUpdateStringRef.current = JSON.stringify(parsed); - - // We have to ensure certain things are always set - try { - // parsed.config.process[0].type = 'ui_trainer'; - parsed.config.process[0].sqlite_db_path = './aitk_db.db'; - parsed.config.process[0].training_folder = settings.TRAINING_FOLDER; - parsed.config.process[0].device = 'cuda'; - parsed.config.process[0].performance_log_every = 10; - } catch (e) { - console.warn(e); + if (JSON.stringify(parsed) !== lastConfigUpdateStringRef.current) { + if (transformOnParse) { + parsed = transformOnParse(parsed); } - migrateJobConfig(parsed); - setJobConfig(parsed); + lastConfigUpdateStringRef.current = JSON.stringify(parsed); + setConfig(parsed); } } catch (e: any) { setHasError(true); @@ -175,6 +154,7 @@ export default function AdvancedJob({ jobConfig, setJobConfig, settings }: Props defaultLanguage="yaml" value={editorValue} theme={theme === 'dark' ? 'vs-dark' : 'light'} + className="z-0" onChange={handleChange} onMount={handleEditorDidMount} options={{ diff --git a/ui/src/components/CaptionDatasetModal.tsx b/ui/src/components/CaptionDatasetModal.tsx index ebc5340a..9e75d6c1 100644 --- a/ui/src/components/CaptionDatasetModal.tsx +++ b/ui/src/components/CaptionDatasetModal.tsx @@ -2,32 +2,19 @@ import React, { useState, useEffect, useRef } from 'react'; import { Modal } from '@/components/Modal'; import { createGlobalState } from 'react-global-hooks'; import { useFromNull } from '@/hooks/useFromNull'; -import { - Checkbox, - CreatableSelectInput, - FormGroup, - SelectInput, - TextAreaInput, - TextInput, -} from '@/components/formInputs'; import { CaptionJobConfig } from '@/types'; -import { defaultCaptionJobConfig, handleCaptionerTypeChange } from '@/helpers/captionJobConfig'; +import { defaultCaptionJobConfig } from '@/helpers/captionJobConfig'; import { objectCopy } from '@/utils/basic'; import { useNestedState } from '@/utils/hooks'; -import { - captionerTypes, - defaultQtype, - groupedCaptionerTypes, - maxNewTokensOptions, - maxResOptions, - quantizationOptions, -} from '@/helpers/captionOptions'; import { isMac } from '@/helpers/basic'; import useGPUInfo from '@/hooks/useGPUInfo'; import { apiClient } from '@/utils/api'; import { v4 as uuidv4 } from 'uuid'; import { startJob } from '@/utils/jobs'; import { startQueue } from '@/utils/queue'; +import CaptionSimpleJob from '@/components/CaptionSimpleJob'; +import AdvancedConfigEditor from '@/components/AdvancedConfigEditor'; +import { SelectInput } from '@/components/formInputs'; export interface CaptionDatasetModalState { datasetPath: string; @@ -45,6 +32,7 @@ export const CaptionDatasetModal: React.FC = () => { const [jobConfig, setJobConfig] = useNestedState(objectCopy(defaultCaptionJobConfig)); const [gpuIDs, setGpuIDs] = useState(null); const { gpuList, isGPUInfoLoaded } = useGPUInfo(); + const [activeTab, setActiveTab] = useState<'simple' | 'advanced'>('simple'); const open = modalInfo !== null; const isSavingRef = useRef(false); const showGPUSelect = !isMac(); @@ -52,6 +40,7 @@ export const CaptionDatasetModal: React.FC = () => { useFromNull(() => { // reset the state setJobConfig(objectCopy(defaultCaptionJobConfig)); + setActiveTab('simple'); // set the path_to_caption if (modalInfo?.datasetPath) { setJobConfig(modalInfo.datasetPath, 'config.process[0].caption.path_to_caption'); @@ -73,8 +62,6 @@ export const CaptionDatasetModal: React.FC = () => { setModalInfo(null); }; - const selectedCaptionOption = captionerTypes.find(option => option.name === jobConfig.config.process[0].type); - const saveJob = async () => { if (isSavingRef.current) return; if (!modalInfo?.datasetPath) { @@ -111,163 +98,54 @@ export const CaptionDatasetModal: React.FC = () => { }); }; - const additionalSections = selectedCaptionOption?.additionalSections || []; + const tabButtonClass = (tab: 'simple' | 'advanced') => + `px-4 py-2 text-sm font-medium border-b-2 transition-colors ${ + activeTab === tab + ? 'border-blue-500 text-blue-400' + : 'border-transparent text-gray-400 hover:text-gray-200 hover:border-gray-600' + }`; return ( - +
+
+ + +
+ {activeTab === 'advanced' && showGPUSelect && ( +
+ setGpuIDs(value)} + options={gpuList.map((gpu: any) => ({ value: `${gpu.index}`, label: `GPU #${gpu.index}` }))} + /> +
+ )} +
{ e.preventDefault(); saveJob(); }} > -
-
-
- { - handleCaptionerTypeChange(jobConfig.config.process[0].type, value, jobConfig, setJobConfig); - }} - options={groupedCaptionerTypes} - /> -
- {showGPUSelect && ( -
- setGpuIDs(value)} - options={gpuList.map((gpu: any) => ({ value: `${gpu.index}`, label: `GPU #${gpu.index}` }))} - /> -
- )} + {activeTab === 'simple' ? ( + + ) : ( +
+
-
- { - if (value?.trim() === '') { - value = null; - } - setJobConfig(value, 'config.process[0].caption.model_name_or_path'); - }} - placeholder="" - options={selectedCaptionOption?.name_or_path_options || []} - required - /> -
- {additionalSections.includes('caption.model_name_or_path2') && ( -
- { - if (value?.trim() === '') { - value = null; - } - setJobConfig(value, 'config.process[0].caption.model_name_or_path2'); - }} - placeholder="" - options={selectedCaptionOption?.name_or_path2_options || []} - /> -
- )} - {additionalSections.includes('caption.fixed_caption') && ( -
- { - if (value?.trim() === '') { - //@ts-ignore - value = undefined; - } - setJobConfig(value, 'config.process[0].caption.fixed_caption'); - }} - placeholder="Enter fixed caption (if you want the same caption for all audio files)" - /> -
- )} -
-
- { - if (value === '') { - setJobConfig(false, 'config.process[0].caption.quantize'); - value = defaultQtype; - } else { - setJobConfig(true, 'config.process[0].caption.quantize'); - } - setJobConfig(value, 'config.process[0].caption.qtype'); - }} - options={quantizationOptions} - /> - {additionalSections.includes('caption.max_res') && ( -
- { - const intVal = parseInt(value); - if (!isNaN(intVal)) { - setJobConfig(intVal, 'config.process[0].caption.max_res'); - } - }} - options={maxResOptions} - /> -
- )} - {additionalSections.includes('caption.max_new_tokens') && ( -
- { - const intVal = parseInt(value); - if (!isNaN(intVal)) { - setJobConfig(intVal, 'config.process[0].caption.max_new_tokens'); - } - }} - options={maxNewTokensOptions} - /> -
- )} -
-
- - setJobConfig(value, 'config.process[0].caption.low_vram')} - /> - setJobConfig(value, 'config.process[0].caption.recaption')} - /> - -
-
- {additionalSections.includes('caption.caption_prompt') && ( -
- { - setJobConfig(value, 'config.process[0].caption.caption_prompt'); - }} - placeholder="Enter caption prompt" - /> -
- )} -
+ )}