Added support for FLUX.1-Kontext-dev

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

@@ -101,10 +101,14 @@ export default function SimpleJob({
                 setJobConfig(value, 'config.process[0].model.arch');
 
                 // update controls for datasets
+                const hasControlPath = newArch?.additionalSections?.includes('datasets.control_path') || false;
                 const controls = newArch?.controls ?? [];
                 const datasets = jobConfig.config.process[0].datasets.map(dataset => {
                   const newDataset = objectCopy(dataset);
                   newDataset.controls = controls;
+                  if (!hasControlPath) {
+                    newDataset.control_path = null; // reset control path if not applicable
+                  }
                   return newDataset;
                 });
                 setJobConfig(datasets, 'config.process[0].datasets');
@@ -412,6 +416,17 @@ export default function SimpleJob({
                         onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].folder_path`)}
                         options={datasetOptions}
                       />
+                      {modelArch?.additionalSections?.includes('datasets.control_path') && (
+                        <SelectInput
+                          label="Control Dataset"
+                          docKey="datasets.control_path"
+                          value={dataset.control_path ?? ''}
+                          onChange={value =>
+                            setJobConfig(value == '' ? null : value, `config.process[0].datasets[${i}].control_path`)
+                          }
+                          options={[{ value: '', label: <>&nbsp;</> }, ...datasetOptions]}
+                        />
+                      )}
                       <NumberInput
                         label="LoRA Weight"
                         value={dataset.network_weight}
@@ -604,6 +619,18 @@ export default function SimpleJob({
               )}
             </div>
             <FormGroup label={`Sample Prompts (${jobConfig.config.process[0].sample.prompts.length})`} className="pt-2">
+              {
+                modelArch?.additionalSections?.includes('sample.ctrl_img') && (
+                  <div className='text-sm text-gray-100 mb-2 py-2 px-4 bg-yellow-700 rounded-lg'>
+                    <p className='font-semibold mb-1'>
+                      Control Images
+                    </p>
+                    To use control images on samples, add --ctrl_img to the prompts below.
+                    <br />
+                    Example: <code className='bg-yellow-900 p-1'>make this a cartoon --ctrl_img /path/to/image.png</code>
+                  </div>
+                )
+              }
               {jobConfig.config.process[0].sample.prompts.map((prompt, i) => (
                 <div key={i} className="flex items-center space-x-2">
                   <div className="flex-1">

ui/src/app/jobs/new/jobConfig.ts

@@ -2,6 +2,7 @@ import { JobConfig, DatasetConfig } from '@/types';
 
 export const defaultDatasetConfig: DatasetConfig = {
   folder_path: '/path/to/images/folder',
+  control_path: null,
   mask_path: null,
   mask_min_value: 0.1,
   default_caption: '',

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

@@ -1,5 +1,8 @@
 type Control = 'depth' | 'line' | 'pose' | 'inpaint';
 
+type DisableableSections = 'model.quantize' | 'train.timestep_type' | 'network.conv';
+type AdditionalSections = 'datasets.control_path' | 'sample.ctrl_img'
+
 export interface ModelArch {
   name: string;
   label: string;
@@ -7,11 +10,11 @@ export interface ModelArch {
   isVideoModel?: boolean;
   defaults?: { [key: string]: any };
   disableSections?: DisableableSections[];
+  additionalSections?: AdditionalSections[];
 }
 
 const defaultNameOrPath = '';
 
-type DisableableSections = 'model.quantize' | 'train.timestep_type' | 'network.conv';
 
 export const modelArchs: ModelArch[] = [
   {
@@ -27,6 +30,21 @@ export const modelArchs: ModelArch[] = [
     },
     disableSections: ['network.conv'],
   },
+  {
+    name: 'flux_kontext',
+    label: 'FLUX.1-Kontext-dev',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['black-forest-labs/FLUX.1-Kontext-dev', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.timestep_type': ['weighted', 'sigmoid'],
+    },
+    disableSections: ['network.conv'],
+    additionalSections: ['datasets.control_path', 'sample.ctrl_img'],
+  },
   {
     name: 'flex1',
     label: 'Flex.1',

ui/src/docs.tsx

@@ -48,6 +48,15 @@ const docs: { [key: string]: ConfigDoc } = {
       </>
     ),
   },
+  'datasets.control_path': {
+    title: 'Control Dataset',
+    description: (
+      <>
+        The control dataset needs to have files that match the filenames of your training dataset. They should be matching file pairs. 
+        These images are fed as control/input images during training. 
+      </>
+    ),
+  },
 };
 
 export const getDoc = (key: string | null | undefined): ConfigDoc | null => {

ui/src/types.ts

@@ -83,6 +83,7 @@ export interface DatasetConfig {
   cache_latents_to_disk?: boolean;
   resolution: number[];
   controls: string[];
+  control_path: string | null;
 }
 
 export interface EMAConfig {