Add ability to train with just music or vocals on yue2. Doesnt work as expected though

ui/src/app/jobs/new/SimpleJob.tsx

@@ -341,16 +341,26 @@ export default function SimpleJob({
                 placeholder=""
               />
             )}
-            {modelArch?.customModelSelectOptions?.map(customOption => (
-              <SelectInput
-                key={customOption.label}
-                label={customOption.label}
-                value={customOption.getValue(jobConfig) ?? ''}
-                doc={customOption.doc}
-                onChange={value => customOption.onChange(value, jobConfig, setJobConfig)}
-                options={customOption.options}
-              />
-            ))}
+            {modelArch?.customModelSelectOptions?.map(customOption =>
+              customOption.type === 'checkbox' ? (
+                <Checkbox
+                  key={customOption.label}
+                  label={customOption.label}
+                  checked={customOption.getValue(jobConfig)}
+                  doc={customOption.doc}
+                  onChange={value => customOption.onChange(value, jobConfig, setJobConfig)}
+                />
+              ) : (
+                <SelectInput
+                  key={customOption.label}
+                  label={customOption.label}
+                  value={customOption.getValue(jobConfig) ?? ''}
+                  doc={customOption.doc}
+                  onChange={value => customOption.onChange(value, jobConfig, setJobConfig)}
+                  options={customOption.options}
+                />
+              ),
+            )}
             {modelArch?.modelNotes && (
               <div className="pt-2">
                 <button

ui/src/app/jobs/new/options.tsx

@@ -49,6 +49,7 @@ type AdditionalSections =
 type ModelGroup = 'image' | 'instruction' | 'video' | 'experimental' | 'audio' | 'llm';
 
 export interface CustomModelSelectOption {
+  type?: 'select';
   label: string;
   options: SelectOption[];
   getValue: (config: JobConfig) => string | undefined;
@@ -56,6 +57,16 @@ export interface CustomModelSelectOption {
   doc?: ConfigDoc;
 }
 
+export interface CustomModelCheckboxOption {
+  type: 'checkbox';
+  label: string;
+  getValue: (config: JobConfig) => boolean;
+  onChange: (value: boolean, config: JobConfig, setJobConfig: (value: any, key: string) => void) => void;
+  doc?: ConfigDoc;
+}
+
+export type CustomModelOption = CustomModelSelectOption | CustomModelCheckboxOption;
+
 export type SampleTag = {
   title: string;
   type: 'text' | 'multiline' | 'number';
@@ -110,7 +121,7 @@ export interface ModelArch {
   sampleTags?: SampleTags;
   gateUrl?: string;
   modelNotes?: React.ReactNode;
-  customModelSelectOptions?: CustomModelSelectOption[];
+  customModelSelectOptions?: CustomModelOption[];
 }
 
 const defaultNameOrPath = '';