Add advanced config section to the captioner

ui/src/components/AdvancedConfigEditor.tsx

@@ -0,0 +1,168 @@
+'use client';
+import { useEffect, useState, useRef } from 'react';
+import YAML from 'yaml';
+import Editor, { OnMount } from '@monaco-editor/react';
+import type { editor } from 'monaco-editor';
+import { useTheme } from '@/components/ThemeProvider';
+
+type Props<T> = {
+  config: T;
+  setConfig: (value: any, key?: string) => void;
+  transformOnParse?: (parsed: any) => any;
+};
+
+const yamlConfig: YAML.DocumentOptions &
+  YAML.SchemaOptions &
+  YAML.ParseOptions &
+  YAML.CreateNodeOptions &
+  YAML.ToStringOptions = {
+  indent: 2,
+  lineWidth: 999999999999,
+  defaultStringType: 'QUOTE_DOUBLE',
+  defaultKeyType: 'PLAIN',
+  directives: true,
+};
+
+function toYaml(obj: any): string {
+  const doc = new YAML.Document(obj, yamlConfig);
+  YAML.visit(doc, {
+    Scalar(_key, node) {
+      if (typeof node.value === 'string' && node.value.includes('\n')) {
+        node.type = YAML.Scalar.BLOCK_LITERAL;
+      }
+    },
+  });
+  return doc.toString(yamlConfig);
+}
+
+export default function AdvancedConfigEditor<T>({ config, setConfig, transformOnParse }: Props<T>) {
+  const { theme } = useTheme();
+  const [editorValue, setEditorValue] = useState<string>('');
+  const [hasError, setHasError] = useState(false);
+  const lastConfigUpdateStringRef = useRef('');
+  const editorRef = useRef<editor.IStandaloneCodeEditor | null>(null);
+  const monacoRef = useRef<any>(null);
+
+  // Track if the editor has been mounted
+  const isEditorMounted = useRef(false);
+
+  // Handler for editor mounting
+  const handleEditorDidMount: OnMount = (editor, monaco) => {
+    editorRef.current = editor;
+    monacoRef.current = monaco;
+    isEditorMounted.current = true;
+
+    // Initial content setup
+    try {
+      const yamlContent = toYaml(config);
+      setEditorValue(yamlContent);
+      lastConfigUpdateStringRef.current = JSON.stringify(config);
+    } catch (e) {
+      console.warn(e);
+    }
+  };
+
+  useEffect(() => {
+    const lastUpdate = lastConfigUpdateStringRef.current;
+    const currentUpdate = JSON.stringify(config);
+
+    // Skip if no changes or editor not yet mounted
+    if (lastUpdate === currentUpdate || !isEditorMounted.current) {
+      return;
+    }
+
+    try {
+      // Preserve cursor position and selection
+      const editor = editorRef.current;
+      if (editor) {
+        // Save current editor state
+        const position = editor.getPosition();
+        const selection = editor.getSelection();
+        const scrollTop = editor.getScrollTop();
+
+        // Update content
+        const yamlContent = toYaml(config);
+
+        // Only update if the content is actually different
+        if (yamlContent !== editor.getValue()) {
+          // Set value directly on the editor model instead of using React state
+          editor.getModel()?.setValue(yamlContent);
+
+          // Restore cursor position and selection
+          if (position) editor.setPosition(position);
+          if (selection) editor.setSelection(selection);
+          editor.setScrollTop(scrollTop);
+        }
+
+        lastConfigUpdateStringRef.current = currentUpdate;
+      }
+    } catch (e) {
+      console.warn(e);
+    }
+  }, [config]);
+
+  const setMarkers = (errors: { message: string; line: number }[]) => {
+    const monaco = monacoRef.current;
+    const model = editorRef.current?.getModel();
+    if (!monaco || !model) return;
+    const markers = errors.map(err => ({
+      severity: monaco.MarkerSeverity.Error,
+      message: err.message,
+      startLineNumber: err.line,
+      startColumn: 1,
+      endLineNumber: err.line,
+      endColumn: model.getLineMaxColumn(err.line),
+    }));
+    monaco.editor.setModelMarkers(model, 'yaml', markers);
+  };
+
+  const handleChange = (value: string | undefined) => {
+    if (value === undefined) return;
+
+    try {
+      let parsed = YAML.parse(value);
+      setHasError(false);
+      setMarkers([]);
+
+      // Don't update config if the change came from the editor itself
+      // to avoid a circular update loop
+      if (JSON.stringify(parsed) !== lastConfigUpdateStringRef.current) {
+        if (transformOnParse) {
+          parsed = transformOnParse(parsed);
+        }
+        lastConfigUpdateStringRef.current = JSON.stringify(parsed);
+        setConfig(parsed);
+      }
+    } catch (e: any) {
+      setHasError(true);
+      const line = e?.linePos?.[0]?.line ?? e?.linePos?.line ?? 1;
+      setMarkers([{ message: e?.message ?? 'Invalid YAML', line }]);
+    }
+  };
+
+  return (
+    <div className="relative h-full w-full">
+      {hasError && (
+        <div
+          className="absolute inset-0 z-10 pointer-events-none rounded-sm"
+          style={{ boxShadow: 'inset 0 0 12px 2px rgba(239, 68, 68, 0.5)' }}
+        />
+      )}
+      <Editor
+        height="100%"
+        width="100%"
+        defaultLanguage="yaml"
+        value={editorValue}
+        theme={theme === 'dark' ? 'vs-dark' : 'light'}
+        className="z-0"
+        onChange={handleChange}
+        onMount={handleEditorDidMount}
+        options={{
+          minimap: { enabled: true },
+          scrollBeyondLastLine: false,
+          automaticLayout: true,
+        }}
+      />
+    </div>
+  );
+}

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<CaptionJobConfig>(objectCopy(defaultCaptionJobConfig));
   const [gpuIDs, setGpuIDs] = useState<string | null>(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 (
-    <Modal isOpen={open} onClose={handleClose} title="Caption Dataset" size="lg">
+    <Modal isOpen={open} onClose={handleClose} title="Caption Dataset" size={activeTab === 'advanced' ? 'xl' : 'lg'}>
       <div className="space-y-4 text-gray-200">
+        <div className="flex items-center border-b border-gray-700 -mt-2">
+          <button type="button" className={tabButtonClass('simple')} onClick={() => setActiveTab('simple')}>
+            Simple
+          </button>
+          <button type="button" className={tabButtonClass('advanced')} onClick={() => setActiveTab('advanced')}>
+            Advanced
+          </button>
+          <div className="flex-1" />
+          {activeTab === 'advanced' && showGPUSelect && (
+            <div className="pb-2">
+              <SelectInput
+                value={`${gpuIDs}`}
+                onChange={value => setGpuIDs(value)}
+                options={gpuList.map((gpu: any) => ({ value: `${gpu.index}`, label: `GPU #${gpu.index}` }))}
+              />
+            </div>
+          )}
+        </div>
         <form
           onSubmit={e => {
             e.preventDefault();
             saveJob();
           }}
         >
-          <div className="text-sm text-gray-400">
-            <div className="grid grid-cols-1 md:grid-cols-2 gap-4 mt-4">
-              <div>
-                <SelectInput
-                  label="Captioner Type"
-                  value={jobConfig.config.process[0].type}
-                  onChange={value => {
-                    handleCaptionerTypeChange(jobConfig.config.process[0].type, value, jobConfig, setJobConfig);
-                  }}
-                  options={groupedCaptionerTypes}
-                />
-              </div>
-              {showGPUSelect && (
-                <div>
-                  <SelectInput
-                    label="GPU ID"
-                    value={`${gpuIDs}`}
-                    onChange={value => setGpuIDs(value)}
-                    options={gpuList.map((gpu: any) => ({ value: `${gpu.index}`, label: `GPU #${gpu.index}` }))}
-                  />
-                </div>
-              )}
+          {activeTab === 'simple' ? (
+            <CaptionSimpleJob
+              jobConfig={jobConfig}
+              setJobConfig={setJobConfig}
+              gpuIDs={gpuIDs}
+              setGpuIDs={setGpuIDs}
+              gpuList={gpuList}
+              showGPUSelect={showGPUSelect}
+            />
+          ) : (
+            <div className="h-[60vh] mt-2">
+              <AdvancedConfigEditor config={jobConfig} setConfig={setJobConfig} />
             </div>
-            <div className="mt-4">
-              <CreatableSelectInput
-                label="Name or Path"
-                value={jobConfig.config.process[0].caption.model_name_or_path}
-                docKey="config.process[0].caption.model_name_or_path"
-                onChange={(value: string | null) => {
-                  if (value?.trim() === '') {
-                    value = null;
-                  }
-                  setJobConfig(value, 'config.process[0].caption.model_name_or_path');
-                }}
-                placeholder=""
-                options={selectedCaptionOption?.name_or_path_options || []}
-                required
-              />
-            </div>
-            {additionalSections.includes('caption.model_name_or_path2') && (
-              <div className="mt-4">
-                <CreatableSelectInput
-                  label="Name or Path 2"
-                  value={jobConfig.config.process[0].caption.model_name_or_path2 || ''}
-                  onChange={(value: string | null) => {
-                    if (value?.trim() === '') {
-                      value = null;
-                    }
-                    setJobConfig(value, 'config.process[0].caption.model_name_or_path2');
-                  }}
-                  placeholder=""
-                  options={selectedCaptionOption?.name_or_path2_options || []}
-                />
-              </div>
-            )}
-            {additionalSections.includes('caption.fixed_caption') && (
-              <div className="mt-4">
-                <TextInput
-                  label="Fixed Caption"
-                  value={jobConfig.config.process[0].caption.fixed_caption || ''}
-                  onChange={value => {
-                    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)"
-                />
-              </div>
-            )}
-            <div className="grid grid-cols-1 md:grid-cols-2 gap-4 mt-4">
-              <div>
-                <SelectInput
-                  label="Quantize"
-                  value={jobConfig.config.process[0].caption.quantize ? jobConfig.config.process[0].caption.qtype : ''}
-                  onChange={value => {
-                    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') && (
-                  <div className="mt-4">
-                    <SelectInput
-                      label="Max Resolution"
-                      value={`${jobConfig.config.process[0].caption.max_res || ''}`}
-                      onChange={value => {
-                        const intVal = parseInt(value);
-                        if (!isNaN(intVal)) {
-                          setJobConfig(intVal, 'config.process[0].caption.max_res');
-                        }
-                      }}
-                      options={maxResOptions}
-                    />
-                  </div>
-                )}
-                {additionalSections.includes('caption.max_new_tokens') && (
-                  <div className="mt-4">
-                    <SelectInput
-                      label="Max New Tokens"
-                      value={`${jobConfig.config.process[0].caption.max_new_tokens || ''}`}
-                      onChange={value => {
-                        const intVal = parseInt(value);
-                        if (!isNaN(intVal)) {
-                          setJobConfig(intVal, 'config.process[0].caption.max_new_tokens');
-                        }
-                      }}
-                      options={maxNewTokensOptions}
-                    />
-                  </div>
-                )}
-              </div>
-              <div>
-                <FormGroup label="Options">
-                  <Checkbox
-                    label="Low VRAM"
-                    checked={jobConfig.config.process[0].caption.low_vram}
-                    onChange={value => setJobConfig(value, 'config.process[0].caption.low_vram')}
-                  />
-                  <Checkbox
-                    label="Recaption"
-                    checked={jobConfig.config.process[0].caption.recaption}
-                    onChange={value => setJobConfig(value, 'config.process[0].caption.recaption')}
-                  />
-                </FormGroup>
-              </div>
-            </div>
-            {additionalSections.includes('caption.caption_prompt') && (
-              <div className="mt-4">
-                <TextAreaInput
-                  label="Caption Prompt"
-                  value={jobConfig.config.process[0].caption.caption_prompt || ''}
-                  onChange={value => {
-                    setJobConfig(value, 'config.process[0].caption.caption_prompt');
-                  }}
-                  placeholder="Enter caption prompt"
-                />
-              </div>
-            )}
-          </div>
+          )}
 
           <div className="mt-6 flex justify-end space-x-3">
             <button

ui/src/components/CaptionSimpleJob.tsx

@@ -0,0 +1,184 @@
+import React from 'react';
+import {
+  Checkbox,
+  CreatableSelectInput,
+  FormGroup,
+  SelectInput,
+  TextAreaInput,
+  TextInput,
+} from '@/components/formInputs';
+import { CaptionJobConfig } from '@/types';
+import { handleCaptionerTypeChange } from '@/helpers/captionJobConfig';
+import {
+  captionerTypes,
+  defaultQtype,
+  groupedCaptionerTypes,
+  maxNewTokensOptions,
+  maxResOptions,
+  quantizationOptions,
+} from '@/helpers/captionOptions';
+
+type Props = {
+  jobConfig: CaptionJobConfig;
+  setJobConfig: (value: any, key?: string) => void;
+  gpuIDs: string | null;
+  setGpuIDs: (value: string | null) => void;
+  gpuList: any;
+  showGPUSelect: boolean;
+};
+
+const CaptionSimpleJob: React.FC<Props> = ({ jobConfig, setJobConfig, gpuIDs, setGpuIDs, gpuList, showGPUSelect }) => {
+  const selectedCaptionOption = captionerTypes.find(option => option.name === jobConfig.config.process[0].type);
+  const additionalSections = selectedCaptionOption?.additionalSections || [];
+
+  return (
+    <div className="text-sm text-gray-400">
+      <div className="grid grid-cols-1 md:grid-cols-2 gap-4 mt-4">
+        <div>
+          <SelectInput
+            label="Captioner Type"
+            value={jobConfig.config.process[0].type}
+            onChange={value => {
+              handleCaptionerTypeChange(jobConfig.config.process[0].type, value, jobConfig, setJobConfig);
+            }}
+            options={groupedCaptionerTypes}
+          />
+        </div>
+        {showGPUSelect && (
+          <div>
+            <SelectInput
+              label="GPU ID"
+              value={`${gpuIDs}`}
+              onChange={value => setGpuIDs(value)}
+              options={gpuList.map((gpu: any) => ({ value: `${gpu.index}`, label: `GPU #${gpu.index}` }))}
+            />
+          </div>
+        )}
+      </div>
+      <div className="mt-4">
+        <CreatableSelectInput
+          label="Name or Path"
+          value={jobConfig.config.process[0].caption.model_name_or_path}
+          docKey="config.process[0].caption.model_name_or_path"
+          onChange={(value: string | null) => {
+            if (value?.trim() === '') {
+              value = null;
+            }
+            setJobConfig(value, 'config.process[0].caption.model_name_or_path');
+          }}
+          placeholder=""
+          options={selectedCaptionOption?.name_or_path_options || []}
+          required
+        />
+      </div>
+      {additionalSections.includes('caption.model_name_or_path2') && (
+        <div className="mt-4">
+          <CreatableSelectInput
+            label="Name or Path 2"
+            value={jobConfig.config.process[0].caption.model_name_or_path2 || ''}
+            onChange={(value: string | null) => {
+              if (value?.trim() === '') {
+                value = null;
+              }
+              setJobConfig(value, 'config.process[0].caption.model_name_or_path2');
+            }}
+            placeholder=""
+            options={selectedCaptionOption?.name_or_path2_options || []}
+          />
+        </div>
+      )}
+      {additionalSections.includes('caption.fixed_caption') && (
+        <div className="mt-4">
+          <TextInput
+            label="Fixed Caption"
+            value={jobConfig.config.process[0].caption.fixed_caption || ''}
+            onChange={value => {
+              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)"
+          />
+        </div>
+      )}
+      <div className="grid grid-cols-1 md:grid-cols-2 gap-4 mt-4">
+        <div>
+          <SelectInput
+            label="Quantize"
+            value={jobConfig.config.process[0].caption.quantize ? jobConfig.config.process[0].caption.qtype : ''}
+            onChange={value => {
+              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') && (
+            <div className="mt-4">
+              <SelectInput
+                label="Max Resolution"
+                value={`${jobConfig.config.process[0].caption.max_res || ''}`}
+                onChange={value => {
+                  const intVal = parseInt(value);
+                  if (!isNaN(intVal)) {
+                    setJobConfig(intVal, 'config.process[0].caption.max_res');
+                  }
+                }}
+                options={maxResOptions}
+              />
+            </div>
+          )}
+          {additionalSections.includes('caption.max_new_tokens') && (
+            <div className="mt-4">
+              <SelectInput
+                label="Max New Tokens"
+                value={`${jobConfig.config.process[0].caption.max_new_tokens || ''}`}
+                onChange={value => {
+                  const intVal = parseInt(value);
+                  if (!isNaN(intVal)) {
+                    setJobConfig(intVal, 'config.process[0].caption.max_new_tokens');
+                  }
+                }}
+                options={maxNewTokensOptions}
+              />
+            </div>
+          )}
+        </div>
+        <div>
+          <FormGroup label="Options">
+            <Checkbox
+              label="Low VRAM"
+              checked={jobConfig.config.process[0].caption.low_vram}
+              onChange={value => setJobConfig(value, 'config.process[0].caption.low_vram')}
+            />
+            <Checkbox
+              label="Recaption"
+              checked={jobConfig.config.process[0].caption.recaption}
+              onChange={value => setJobConfig(value, 'config.process[0].caption.recaption')}
+            />
+          </FormGroup>
+        </div>
+      </div>
+      {additionalSections.includes('caption.caption_prompt') && (
+        <div className="mt-4">
+          <TextAreaInput
+            label="Caption Prompt"
+            value={jobConfig.config.process[0].caption.caption_prompt || ''}
+            onChange={value => {
+              setJobConfig(value, 'config.process[0].caption.caption_prompt');
+            }}
+            placeholder="Enter caption prompt"
+          />
+        </div>
+      )}
+    </div>
+  );
+};
+
+export default CaptionSimpleJob;