From 1069dee0e444c900d2eba78ab6f4ef78a43f3ffc Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Thu, 25 Sep 2025 11:10:02 -0600 Subject: [PATCH] Added ui sopport for multi control samples and datasets. Added qwen image edit 5209 to the ui --- .../train_lora_qwen_image_edit_2509_32gb.yaml | 105 +++++++++ toolkit/config_modules.py | 15 ++ ui/src/app/jobs/new/SimpleJob.tsx | 179 +++++++-------- ui/src/app/jobs/new/jobConfig.ts | 1 - ui/src/app/jobs/new/options.ts | 25 +++ ui/src/app/jobs/new/utils.ts | 105 +++++++++ ui/src/components/SampleControlImage.tsx | 206 ++++++++++++++++++ ui/src/docs.tsx | 16 ++ ui/src/types.ts | 8 +- version.py | 2 +- 10 files changed, 574 insertions(+), 88 deletions(-) create mode 100644 config/examples/train_lora_qwen_image_edit_2509_32gb.yaml create mode 100644 ui/src/app/jobs/new/utils.ts create mode 100644 ui/src/components/SampleControlImage.tsx diff --git a/config/examples/train_lora_qwen_image_edit_2509_32gb.yaml b/config/examples/train_lora_qwen_image_edit_2509_32gb.yaml new file mode 100644 index 00000000..845ac92b --- /dev/null +++ b/config/examples/train_lora_qwen_image_edit_2509_32gb.yaml @@ -0,0 +1,105 @@ +--- +job: extension +config: + # this name will be the folder and filename name + name: "my_first_qwen_image_edit_2509_lora_v1" + process: + - type: 'diffusion_trainer' + # root folder to save training sessions/samples/weights + training_folder: "output" + # uncomment to see performance stats in the terminal every N steps +# performance_log_every: 1000 + device: cuda:0 + network: + type: "lora" + linear: 16 + linear_alpha: 16 + save: + dtype: float16 # precision to save + save_every: 250 # save every this many steps + max_step_saves_to_keep: 4 # how many intermittent saves to keep + datasets: + # datasets are a folder of images. captions need to be txt files with the same name as the image + # for instance image2.jpg and image2.txt. Only jpg, jpeg, and png are supported currently + # images will automatically be resized and bucketed into the resolution specified + # on windows, escape back slashes with another backslash so + # "C:\\path\\to\\images\\folder" + - folder_path: "/path/to/images/folder" + # can do up to 3 control image folders, file names must match target file names, but aspect/size can be different + control_path: + - "/path/to/control/images/folder1" + - "/path/to/control/images/folder2" + - "/path/to/control/images/folder3" + caption_ext: "txt" + # default_caption: "a person" # if caching text embeddings, if you don't have captions, this will get cached + caption_dropout_rate: 0.05 # will drop out the caption 5% of time + resolution: [ 512, 768, 1024 ] # qwen image enjoys multiple resolutions + # a trigger word that can be cached with the text embeddings + # trigger_word: "optional trigger word" + train: + batch_size: 1 + # caching text embeddings is required for 32GB + cache_text_embeddings: true + # unload_text_encoder: true + + steps: 3000 # total number of steps to train 500 - 4000 is a good range + gradient_accumulation: 1 + timestep_type: "weighted" + train_unet: true + train_text_encoder: false # probably won't work with qwen image + gradient_checkpointing: true # need the on unless you have a ton of vram + noise_scheduler: "flowmatch" # for training only + optimizer: "adamw8bit" + lr: 1e-4 + # uncomment this to skip the pre training sample + # skip_first_sample: true + # uncomment to completely disable sampling + # disable_sampling: true + dtype: bf16 + model: + # huggingface model name or path + name_or_path: "Qwen/Qwen-Image-Edit-2509" + arch: "qwen_image_edit_plus" + quantize: true + # to use the ARA use the | pipe to point to hf path, or a local path if you have one. + # 3bit is required for 32GB + qtype: "uint3|ostris/accuracy_recovery_adapters/qwen_image_edit_2509_torchao_uint3.safetensors" + quantize_te: true + qtype_te: "qfloat8" + low_vram: true + sample: + sampler: "flowmatch" # must match train.noise_scheduler + sample_every: 250 # sample every this many steps + width: 1024 + height: 1024 + # you can provide up to 3 control images here + samples: + - prompt: "Do whatever with Image1 and Image2" + ctrl_img_1: "/path/to/image1.png" + ctrl_img_2: "/path/to/image2.png" + # ctrl_img_3: "/path/to/image3.png" + - prompt: "Do whatever with Image1 and Image2" + ctrl_img_1: "/path/to/image1.png" + ctrl_img_2: "/path/to/image2.png" + # ctrl_img_3: "/path/to/image3.png" + - prompt: "Do whatever with Image1 and Image2" + ctrl_img_1: "/path/to/image1.png" + ctrl_img_2: "/path/to/image2.png" + # ctrl_img_3: "/path/to/image3.png" + - prompt: "Do whatever with Image1 and Image2" + ctrl_img_1: "/path/to/image1.png" + ctrl_img_2: "/path/to/image2.png" + # ctrl_img_3: "/path/to/image3.png" + - prompt: "Do whatever with Image1 and Image2" + ctrl_img_1: "/path/to/image1.png" + ctrl_img_2: "/path/to/image2.png" + # ctrl_img_3: "/path/to/image3.png" + neg: "" + seed: 42 + walk_seed: true + guidance_scale: 3 + sample_steps: 25 +# you can add any additional meta info here. [name] is replaced with config name at top +meta: + name: "[name]" + version: '1.0' diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index 24da74fe..ad6cc1f3 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -831,6 +831,21 @@ class DatasetConfig: if self.control_path == '': self.control_path = None + # handle multi control inputs from the ui. It is just easier to handle it here for a cleaner ui experience + control_path_1 = kwargs.get('control_path_1', None) + control_path_2 = kwargs.get('control_path_2', None) + control_path_3 = kwargs.get('control_path_3', None) + + if any([control_path_1, control_path_2, control_path_3]): + control_paths = [] + if control_path_1: + control_paths.append(control_path_1) + if control_path_2: + control_paths.append(control_path_2) + if control_path_3: + control_paths.append(control_path_3) + self.control_path = control_paths + # color for transparent reigon of control images with transparency self.control_transparent_color: List[int] = kwargs.get('control_transparent_color', [0, 0, 0]) # inpaint images should be webp/png images with alpha channel. The alpha 0 (invisible) section will diff --git a/ui/src/app/jobs/new/SimpleJob.tsx b/ui/src/app/jobs/new/SimpleJob.tsx index 3953c182..4ae5bcc9 100644 --- a/ui/src/app/jobs/new/SimpleJob.tsx +++ b/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({ )} - { - 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') && ( + { + setJobConfig(value, 'config.process[0].train.unload_text_encoder'); + if (value) { + setJobConfig(false, 'config.process[0].train.cache_text_embeddings'); + } + }} + /> + )}
setJobConfig(value, `config.process[0].datasets[${i}].folder_path`)} options={datasetOptions} @@ -659,6 +612,49 @@ export default function SimpleJob({ options={[{ value: '', label: <>  }, ...datasetOptions]} /> )} + {modelArch?.additionalSections?.includes('datasets.multi_control_paths') && ( + <> + + setJobConfig( + value == '' ? null : value, + `config.process[0].datasets[${i}].control_path_1`, + ) + } + options={[{ value: '', label: <>  }, ...datasetOptions]} + /> + + setJobConfig( + value == '' ? null : value, + `config.process[0].datasets[${i}].control_path_2`, + ) + } + options={[{ value: '', label: <>  }, ...datasetOptions]} + /> + + setJobConfig( + value == '' ? null : value, + `config.process[0].datasets[${i}].control_path_3`, + ) + } + options={[{ value: '', label: <>  }, ...datasetOptions]} + /> + + )}
- + {modelArch?.additionalSections?.includes('datasets.multi_control_paths') && ( + +
+ {['ctrl_img_1', 'ctrl_img_2', 'ctrl_img_3'].map((ctrlKey, ctrl_idx) => ( + { + 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}`); + } + }} + /> + ))} +
+ + )} {modelArch?.additionalSections?.includes('sample.ctrl_img') && ( -
{ - openAddImageModal(imagePath => { - console.log('Selected image path:', imagePath); - if (!imagePath) return; + { + 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 && ( -
Add Control Image
- )} -
+ /> )}
diff --git a/ui/src/app/jobs/new/jobConfig.ts b/ui/src/app/jobs/new/jobConfig.ts index c18f7327..54f34ca9 100644 --- a/ui/src/app/jobs/new/jobConfig.ts +++ b/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: '', diff --git a/ui/src/app/jobs/new/options.ts b/ui/src/app/jobs/new/options.ts index 7735ae35..eec57717 100644 --- a/ui/src/app/jobs/new/options.ts +++ b/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', diff --git a/ui/src/app/jobs/new/utils.ts b/ui/src/app/jobs/new/utils.ts new file mode 100644 index 00000000..5d3803e8 --- /dev/null +++ b/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'); +}; diff --git a/ui/src/components/SampleControlImage.tsx b/ui/src/components/SampleControlImage.tsx new file mode 100644 index 00000000..b10b2cef --- /dev/null +++ b/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(null); + const fileInputRef = useRef(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 ( +
!isUploading && fileInputRef.current?.click()} + > + {/* Hidden input for click-to-open */} + { + const file = e.currentTarget.files?.[0]; + if (file) handleUpload(file); + }} + /> + + {/* Empty state — centered */} + {!backgroundUrl && ( +
+ +
{instruction}
+
Click or drop
+
+ )} + + {/* Existing image overlays */} + {backgroundUrl && !isUploading && ( + <> +
+
+ + Replace +
+
+ + {/* Clear (X) button */} + + + )} + + {/* Uploading overlay */} + {isUploading && ( +
+
+
+
+
+
Uploading… {uploadProgress}%
+
+
+ )} +
+ ); +} diff --git a/ui/src/docs.tsx b/ui/src/docs.tsx index cadf0e01..b93ae3a2 100644 --- a/ui/src/docs.tsx +++ b/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. +
+
+ 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. ), }, diff --git a/ui/src/types.ts b/ui/src/types.ts index e5ffb25b..71048ca5 100644 --- a/ui/src/types.ts +++ b/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 { diff --git a/version.py b/version.py index dd1e6dbf..af9f7a48 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.5.10" \ No newline at end of file +VERSION = "0.6.0" \ No newline at end of file