Added ui sopport for multi control samples and datasets. Added qwen image edit 5209 to the ui

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

@@ -15,7 +15,9 @@ import { TextInput, SelectInput, Checkbox, FormGroup, NumberInput } from '@/comp
 import Card from '@/components/Card';
 import { X } from 'lucide-react';
 import AddSingleImageModal, { openAddImageModal } from '@/components/AddSingleImageModal';
+import SampleControlImage from '@/components/SampleControlImage';
 import { FlipHorizontal2, FlipVertical2 } from 'lucide-react';
+import { handleModelArchChange } from './utils';
 
 type Props = {
   jobConfig: JobConfig;
@@ -185,58 +187,7 @@ export default function SimpleJob({
               label="Model Architecture"
               value={jobConfig.config.process[0].model.arch}
               onChange={value => {
-                const currentArch = modelArchs.find(a => a.name === jobConfig.config.process[0].model.arch);
-                if (!currentArch || currentArch.name === value) {
-                  return;
-                }
-                // update the defaults when a model is selected
-                const newArch = modelArchs.find(model => model.name === value);
-
-                // update vram setting
-                if (!newArch?.additionalSections?.includes('model.low_vram')) {
-                  setJobConfig(false, 'config.process[0].model.low_vram');
-                }
-
-                // revert defaults from previous model
-                for (const key in currentArch.defaults) {
-                  setJobConfig(currentArch.defaults[key][1], key);
-                }
-
-                if (newArch?.defaults) {
-                  for (const key in newArch.defaults) {
-                    setJobConfig(newArch.defaults[key][0], key);
-                  }
-                }
-                // set new model
-                setJobConfig(value, 'config.process[0].model.arch');
-
-                // update datasets
-                const hasControlPath = newArch?.additionalSections?.includes('datasets.control_path') || false;
-                const hasNumFrames = newArch?.additionalSections?.includes('datasets.num_frames') || 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
-                  }
-                  if (!hasNumFrames) {
-                    newDataset.num_frames = 1; // reset num_frames if not applicable
-                  }
-                  return newDataset;
-                });
-                setJobConfig(datasets, 'config.process[0].datasets');
-
-                // update samples
-                const hasSampleCtrlImg = newArch?.additionalSections?.includes('sample.ctrl_img') || false;
-                const samples = jobConfig.config.process[0].sample.samples.map(sample => {
-                  const newSample = objectCopy(sample);
-                  if (!hasSampleCtrlImg) {
-                    delete newSample.ctrl_img; // remove ctrl_img if not applicable
-                  }
-                  return newSample;
-                });
-                setJobConfig(samples, 'config.process[0].sample.samples');
+                handleModelArchChange(jobConfig.config.process[0].model.arch, value, jobConfig, setJobConfig);
               }}
               options={groupedModelOptions}
             />
@@ -557,17 +508,19 @@ export default function SimpleJob({
                 )}
 
                 <FormGroup label="Text Encoder Optimizations" className="pt-2">
-                  <Checkbox
-                    label="Unload TE"
-                    checked={jobConfig.config.process[0].train.unload_text_encoder || false}
-                    docKey={'train.unload_text_encoder'}
-                    onChange={value => {
-                      setJobConfig(value, 'config.process[0].train.unload_text_encoder');
-                      if (value) {
-                        setJobConfig(false, 'config.process[0].train.cache_text_embeddings');
-                      }
-                    }}
-                  />
+                  {!disableSections.includes('train.unload_text_encoder') && (
+                    <Checkbox
+                      label="Unload TE"
+                      checked={jobConfig.config.process[0].train.unload_text_encoder || false}
+                      docKey={'train.unload_text_encoder'}
+                      onChange={value => {
+                        setJobConfig(value, 'config.process[0].train.unload_text_encoder');
+                        if (value) {
+                          setJobConfig(false, 'config.process[0].train.cache_text_embeddings');
+                        }
+                      }}
+                    />
+                  )}
                   <Checkbox
                     label="Cache Text Embeddings"
                     checked={jobConfig.config.process[0].train.cache_text_embeddings || false}
@@ -642,7 +595,7 @@ export default function SimpleJob({
                   <div className="grid grid-cols-1 md:grid-cols-2 lg:grid-cols-4 gap-6">
                     <div>
                       <SelectInput
-                        label="Dataset"
+                        label="Target Dataset"
                         value={dataset.folder_path}
                         onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].folder_path`)}
                         options={datasetOptions}
@@ -659,6 +612,49 @@ export default function SimpleJob({
                           options={[{ value: '', label: <>&nbsp;</> }, ...datasetOptions]}
                         />
                       )}
+                      {modelArch?.additionalSections?.includes('datasets.multi_control_paths') && (
+                        <>
+                          <SelectInput
+                            label="Control Dataset 1"
+                            docKey="datasets.multi_control_paths"
+                            value={dataset.control_path_1 ?? ''}
+                            className="pt-2"
+                            onChange={value =>
+                              setJobConfig(
+                                value == '' ? null : value,
+                                `config.process[0].datasets[${i}].control_path_1`,
+                              )
+                            }
+                            options={[{ value: '', label: <>&nbsp;</> }, ...datasetOptions]}
+                          />
+                          <SelectInput
+                            label="Control Dataset 2"
+                            docKey="datasets.multi_control_paths"
+                            value={dataset.control_path_2 ?? ''}
+                            className="pt-2"
+                            onChange={value =>
+                              setJobConfig(
+                                value == '' ? null : value,
+                                `config.process[0].datasets[${i}].control_path_2`,
+                              )
+                            }
+                            options={[{ value: '', label: <>&nbsp;</> }, ...datasetOptions]}
+                          />
+                          <SelectInput
+                            label="Control Dataset 3"
+                            docKey="datasets.multi_control_paths"
+                            value={dataset.control_path_3 ?? ''}
+                            className="pt-2"
+                            onChange={value =>
+                              setJobConfig(
+                                value == '' ? null : value,
+                                `config.process[0].datasets[${i}].control_path_3`,
+                              )
+                            }
+                            options={[{ value: '', label: <>&nbsp;</> }, ...datasetOptions]}
+                          />
+                        </>
+                      )}
                       <NumberInput
                         label="LoRA Weight"
                         value={dataset.network_weight}
@@ -1062,30 +1058,43 @@ export default function SimpleJob({
                           />
                         </div>
                       </div>
-
+                      {modelArch?.additionalSections?.includes('datasets.multi_control_paths') && (
+                        <FormGroup label="Control Images" className="pt-2 ml-4">
+                          <div className="grid grid-cols-1 md:grid-cols-3 gap-2 mt-2 mt-2">
+                            {['ctrl_img_1', 'ctrl_img_2', 'ctrl_img_3'].map((ctrlKey, ctrl_idx) => (
+                              <SampleControlImage
+                                key={ctrlKey}
+                                instruction={`Add Control Image ${ctrl_idx + 1}`}
+                                className=""
+                                src={sample[ctrlKey as keyof typeof sample] as string}
+                                onNewImageSelected={imagePath => {
+                                  if (!imagePath) {
+                                    let newSamples = objectCopy(jobConfig.config.process[0].sample.samples);
+                                    delete newSamples[i][ctrlKey as keyof typeof sample];
+                                    setJobConfig(newSamples, 'config.process[0].sample.samples');
+                                  } else {
+                                    setJobConfig(imagePath, `config.process[0].sample.samples[${i}].${ctrlKey}`);
+                                  }
+                                }}
+                              />
+                            ))}
+                          </div>
+                        </FormGroup>
+                      )}
                       {modelArch?.additionalSections?.includes('sample.ctrl_img') && (
-                        <div
-                          className="h-14 w-14 mt-2 ml-4 border border-gray-500 flex items-center justify-center rounded cursor-pointer hover:bg-gray-700 transition-colors"
-                          style={{
-                            backgroundImage: sample.ctrl_img
-                              ? `url(${`/api/img/${encodeURIComponent(sample.ctrl_img)}`})`
-                              : 'none',
-                            backgroundSize: 'cover',
-                            backgroundPosition: 'center',
-                            marginBottom: '-1rem',
-                          }}
-                          onClick={() => {
-                            openAddImageModal(imagePath => {
-                              console.log('Selected image path:', imagePath);
-                              if (!imagePath) return;
+                        <SampleControlImage
+                          className="mt-6 ml-4"
+                          src={sample.ctrl_img}
+                          onNewImageSelected={imagePath => {
+                            if (!imagePath) {
+                              let newSamples = objectCopy(jobConfig.config.process[0].sample.samples);
+                              delete newSamples[i].ctrl_img;
+                              setJobConfig(newSamples, 'config.process[0].sample.samples');
+                            } else {
                               setJobConfig(imagePath, `config.process[0].sample.samples[${i}].ctrl_img`);
-                            });
+                            }
                           }}
-                        >
-                          {!sample.ctrl_img && (
-                            <div className="text-gray-400 text-xs text-center font-bold">Add Control Image</div>
-                          )}
-                        </div>
+                        />
                       )}
                     </div>
                     <div className="pb-4"></div>

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

@@ -2,7 +2,6 @@ import { JobConfig, DatasetConfig, SliderConfig } 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

@@ -9,12 +9,15 @@ type DisableableSections =
   | 'network.conv'
   | 'trigger_word'
   | 'train.diff_output_preservation'
+  | 'train.unload_text_encoder'
   | 'slider';
 
 type AdditionalSections =
   | 'datasets.control_path'
+  | 'datasets.multi_control_paths'
   | 'datasets.do_i2v'
   | 'sample.ctrl_img'
+  | 'sample.multi_ctrl_imgs'
   | 'datasets.num_frames'
   | 'model.multistage'
   | 'model.low_vram';
@@ -335,6 +338,28 @@ export const modelArchs: ModelArch[] = [
       '3 bit with ARA': 'uint3|ostris/accuracy_recovery_adapters/qwen_image_edit_torchao_uint3.safetensors',
     },
   },
+  {
+    name: 'qwen_image_edit_plus',
+    label: 'Qwen-Image-Edit-2509',
+    group: 'instruction',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['Qwen/Qwen-Image-Edit-2509', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].model.low_vram': [true, false],
+      'config.process[0].train.unload_text_encoder': [false, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.timestep_type': ['weighted', 'sigmoid'],
+      'config.process[0].model.qtype': ['qfloat8', 'qfloat8'],
+    },
+    disableSections: ['network.conv', 'train.unload_text_encoder'],
+    additionalSections: ['datasets.multi_control_paths', 'sample.multi_ctrl_imgs', 'model.low_vram'],
+    accuracyRecoveryAdapters: {
+      '3 bit with ARA': 'uint3|ostris/accuracy_recovery_adapters/qwen_image_edit_2509_torchao_uint3.safetensors',
+    },
+  },
   {
     name: 'hidream',
     label: 'HiDream',

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

@@ -0,0 +1,105 @@
+import { GroupedSelectOption, JobConfig, SelectOption } from '@/types';
+import { modelArchs, ModelArch } from './options';
+import { objectCopy } from '@/utils/basic';
+
+export const handleModelArchChange = (
+  currentArchName: string,
+  newArchName: string,
+  jobConfig: JobConfig,
+  setJobConfig: (value: any, key: string) => void,
+) => {
+  const currentArch = modelArchs.find(a => a.name === currentArchName);
+  if (!currentArch || currentArch.name === newArchName) {
+    return;
+  }
+
+  // update the defaults when a model is selected
+  const newArch = modelArchs.find(model => model.name === newArchName);
+
+  // update vram setting
+  if (!newArch?.additionalSections?.includes('model.low_vram')) {
+    setJobConfig(false, 'config.process[0].model.low_vram');
+  }
+
+  // revert defaults from previous model
+  for (const key in currentArch.defaults) {
+    setJobConfig(currentArch.defaults[key][1], key);
+  }
+
+  if (newArch?.defaults) {
+    for (const key in newArch.defaults) {
+      setJobConfig(newArch.defaults[key][0], key);
+    }
+  }
+  // set new model
+  setJobConfig(newArchName, 'config.process[0].model.arch');
+
+  // update datasets
+  const hasControlPath = newArch?.additionalSections?.includes('datasets.control_path') || false;
+  const hasMultiControlPaths = newArch?.additionalSections?.includes('datasets.multi_control_paths') || false;
+  const hasNumFrames = newArch?.additionalSections?.includes('datasets.num_frames') || false;
+  const controls = newArch?.controls ?? [];
+  const datasets = jobConfig.config.process[0].datasets.map(dataset => {
+    const newDataset = objectCopy(dataset);
+    newDataset.controls = controls;
+    if (hasMultiControlPaths) {
+      // make sure the config has the multi control paths
+      newDataset.control_path_1 = newDataset.control_path_1 || null;
+      newDataset.control_path_2 = newDataset.control_path_2 || null;
+      newDataset.control_path_3 = newDataset.control_path_3 || null;
+      // if we previously had a single control path and now
+      // we selected a multi control model
+      if (newDataset.control_path && newDataset.control_path !== '') {
+        // only set if not overwriting
+        if (!newDataset.control_path_1) {
+          newDataset.control_path_1 = newDataset.control_path;
+        }
+      }
+      delete newDataset.control_path; // remove single control path
+    } else if (hasControlPath) {
+      newDataset.control_path = newDataset.control_path || null;
+      if (newDataset.control_path_1 && newDataset.control_path_1 !== '') {
+        newDataset.control_path = newDataset.control_path_1;
+      }
+      if (newDataset.control_path_1) {
+        delete newDataset.control_path_1;
+      }
+      if (newDataset.control_path_2) {
+        delete newDataset.control_path_2;
+      }
+      if (newDataset.control_path_3) {
+        delete newDataset.control_path_3;
+      }
+    } else {
+      // does not have control images
+      if (newDataset.control_path) {
+        delete newDataset.control_path;
+      }
+      if (newDataset.control_path_1) {
+        delete newDataset.control_path_1;
+      }
+      if (newDataset.control_path_2) {
+        delete newDataset.control_path_2;
+      }
+      if (newDataset.control_path_3) {
+        delete newDataset.control_path_3;
+      }
+    }
+    if (!hasNumFrames) {
+      newDataset.num_frames = 1; // reset num_frames if not applicable
+    }
+    return newDataset;
+  });
+  setJobConfig(datasets, 'config.process[0].datasets');
+
+  // update samples
+  const hasSampleCtrlImg = newArch?.additionalSections?.includes('sample.ctrl_img') || false;
+  const samples = jobConfig.config.process[0].sample.samples.map(sample => {
+    const newSample = objectCopy(sample);
+    if (!hasSampleCtrlImg) {
+      delete newSample.ctrl_img; // remove ctrl_img if not applicable
+    }
+    return newSample;
+  });
+  setJobConfig(samples, 'config.process[0].sample.samples');
+};

ui/src/components/SampleControlImage.tsx

@@ -0,0 +1,206 @@
+'use client';
+
+import React, { useCallback, useMemo, useRef, useState } from 'react';
+import classNames from 'classnames';
+import { useDropzone } from 'react-dropzone';
+import { FaUpload, FaImage, FaTimes } from 'react-icons/fa';
+import { apiClient } from '@/utils/api';
+import type { AxiosProgressEvent } from 'axios';
+
+interface Props {
+  src: string | null | undefined;
+  className?: string;
+  instruction?: string;
+  onNewImageSelected: (imagePath: string | null) => void;
+}
+
+export default function SampleControlImage({
+  src,
+  className,
+  instruction = 'Add Control Image',
+  onNewImageSelected,
+}: Props) {
+  const [isUploading, setIsUploading] = useState(false);
+  const [uploadProgress, setUploadProgress] = useState(0);
+  const [localPreview, setLocalPreview] = useState<string | null>(null);
+  const fileInputRef = useRef<HTMLInputElement | null>(null);
+
+  const backgroundUrl = useMemo(() => {
+    if (localPreview) return localPreview;
+    if (src) return `/api/img/${encodeURIComponent(src)}`;
+    return null;
+  }, [src, localPreview]);
+
+  const handleUpload = useCallback(
+    async (file: File) => {
+      if (!file) return;
+      setIsUploading(true);
+      setUploadProgress(0);
+
+      const objectUrl = URL.createObjectURL(file);
+      setLocalPreview(objectUrl);
+
+      const formData = new FormData();
+      formData.append('files', file);
+
+      try {
+        const resp = await apiClient.post(`/api/img/upload`, formData, {
+          headers: { 'Content-Type': 'multipart/form-data' },
+          onUploadProgress: (evt: AxiosProgressEvent) => {
+            const total = evt.total ?? 100;
+            const loaded = evt.loaded ?? 0;
+            setUploadProgress(Math.round((loaded * 100) / total));
+          },
+          timeout: 0,
+        });
+
+        const uploaded = resp?.data?.files?.[0] ?? null;
+        onNewImageSelected(uploaded);
+      } catch (err) {
+        console.error('Upload failed:', err);
+        setLocalPreview(null);
+      } finally {
+        setIsUploading(false);
+        setUploadProgress(0);
+        URL.revokeObjectURL(objectUrl);
+        if (fileInputRef.current) fileInputRef.current.value = '';
+      }
+    },
+    [onNewImageSelected],
+  );
+
+  const onDrop = useCallback(
+    (acceptedFiles: File[]) => {
+      if (acceptedFiles.length === 0) return;
+      handleUpload(acceptedFiles[0]);
+    },
+    [handleUpload],
+  );
+
+  const clearImage = useCallback(
+    (e?: React.MouseEvent) => {
+      console.log('clearImage');
+      if (e) {
+        e.stopPropagation();
+        e.preventDefault();
+      }
+      setLocalPreview(null);
+      onNewImageSelected(null);
+      if (fileInputRef.current) fileInputRef.current.value = '';
+    },
+    [onNewImageSelected],
+  );
+
+  // Drag & drop only; click handled via our own hidden input
+  const { getRootProps, isDragActive } = useDropzone({
+    onDrop,
+    accept: { 'image/*': ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp'] },
+    multiple: false,
+    noClick: true,
+    noKeyboard: true,
+  });
+
+  const rootProps = getRootProps();
+
+  return (
+    <div
+      {...rootProps}
+      className={classNames(
+        'group relative flex items-center justify-center rounded-xl cursor-pointer ring-1 ring-inset',
+        'transition-all duration-200 select-none overflow-hidden text-center',
+        'h-20 w-20',
+        backgroundUrl ? 'bg-gray-800 ring-gray-700' : 'bg-gradient-to-b from-gray-800 to-gray-900 ring-gray-700',
+        isDragActive ? 'outline outline-2 outline-blue-500' : 'hover:ring-gray-600',
+        className,
+      )}
+      style={
+        backgroundUrl
+          ? {
+              backgroundImage: `url("${backgroundUrl}")`,
+              backgroundSize: 'cover',
+              backgroundPosition: 'center',
+            }
+          : undefined
+      }
+      onClick={() => !isUploading && fileInputRef.current?.click()}
+    >
+      {/* Hidden input for click-to-open */}
+      <input
+        ref={fileInputRef}
+        type="file"
+        accept="image/*"
+        className="hidden"
+        onChange={e => {
+          const file = e.currentTarget.files?.[0];
+          if (file) handleUpload(file);
+        }}
+      />
+
+      {/* Empty state — centered */}
+      {!backgroundUrl && (
+        <div className="flex flex-col items-center justify-center text-gray-300 text-center">
+          <FaImage className="opacity-80" />
+          <div className="mt-1 text-[10px] font-semibold tracking-wide opacity-80">{instruction}</div>
+          <div className="mt-0.5 text-[9px] opacity-60">Click or drop</div>
+        </div>
+      )}
+
+      {/* Existing image overlays */}
+      {backgroundUrl && !isUploading && (
+        <>
+          <div
+            className={classNames(
+              'pointer-events-none absolute inset-0 flex items-center justify-center',
+              'bg-black/0 group-hover:bg-black/20',
+              isDragActive && 'bg-black/35',
+              'transition-colors',
+            )}
+          >
+            <div
+              className={classNames(
+                'inline-flex items-center gap-1 rounded-md px-2 py-1',
+                'text-[10px] font-semibold',
+                'bg-black/45 text-white/90 backdrop-blur-sm',
+                'opacity-0 group-hover:opacity-100 transition-opacity',
+              )}
+            >
+              <FaUpload className="text-[10px]" />
+              <span>Replace</span>
+            </div>
+          </div>
+
+          {/* Clear (X) button */}
+          <button
+            type="button"
+            onClick={clearImage}
+            title="Clear image"
+            aria-label="Clear image"
+            className={classNames(
+              'absolute right-1.5 top-1.5 z-10 inline-flex items-center justify-center',
+              'h-5 w-5 rounded-md bg-black/55 text-white/90',
+              'opacity-0 group-hover:opacity-100 transition-opacity',
+              'hover:bg-black/70',
+            )}
+          >
+            <FaTimes className="text-[10px]" />
+          </button>
+        </>
+      )}
+
+      {/* Uploading overlay */}
+      {isUploading && (
+        <div className="absolute inset-0 flex flex-col items-center justify-center bg-black/60 backdrop-blur-[1px] text-center">
+          <div className="w-4/5 max-w-40">
+            <div className="h-1.5 w-full rounded-full bg-white/15">
+              <div
+                className="h-1.5 rounded-full bg-white/80 transition-[width]"
+                style={{ width: `${uploadProgress}%` }}
+              />
+            </div>
+            <div className="mt-1 text-[10px] font-medium text-white/90">Uploading… {uploadProgress}%</div>
+          </div>
+        </div>
+      )}
+    </div>
+  );
+}

ui/src/docs.tsx

@@ -53,10 +53,26 @@ 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. The control images will be 
+        resized to match the training images.
+      </>
+    ),
+  },
+  'datasets.multi_control_paths': {
+    title: 'Multi 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.
+        <br />
+        <br />
+        For multi control datasets, the controls will all be applied in the order they are listed. If the model does not
+        require the images to be the same aspect ratios, such as with Qwen/Qwen-Image-Edit-2509, then the control images
+        do not need to match the aspect size or aspect ratio of the target image and they will be automatically resized to 
+        the ideal resolutions for the model / target images.
       </>
     ),
   },

ui/src/types.ts

@@ -83,12 +83,15 @@ export interface DatasetConfig {
   cache_latents_to_disk?: boolean;
   resolution: number[];
   controls: string[];
-  control_path: string | null;
+  control_path?: string | null;
   num_frames: number;
   shrink_video_to_frames: boolean;
   do_i2v: boolean;
   flip_x: boolean;
   flip_y: boolean;
+  control_path_1?: string | null;
+  control_path_2?: string | null;
+  control_path_3?: string | null;
 }
 
 export interface EMAConfig {
@@ -155,6 +158,9 @@ export interface SampleItem {
   ctrl_img?: string | null;
   ctrl_idx?: number;
   network_multiplier?: number;
+  ctrl_img_1?: string | null;
+  ctrl_img_2?: string | null;
+  ctrl_img_3?: string | null;
 }
 
 export interface SampleConfig {