From 59ff4efae5e3050d1d06ba9becb79edcdba59def Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Fri, 22 Aug 2025 10:26:46 -0600
Subject: [PATCH] Add support for training  Qwen Image Edit in the UI

---
 ui/src/app/jobs/new/options.ts | 27 ++++++++++++++++++++++++---
 version.py                     |  2 +-
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/ui/src/app/jobs/new/options.ts b/ui/src/app/jobs/new/options.ts
index 1381495d..2c33ad72 100644
--- a/ui/src/app/jobs/new/options.ts
+++ b/ui/src/app/jobs/new/options.ts
@@ -10,7 +10,7 @@ type AdditionalSections =
   | 'datasets.num_frames'
   | 'model.multistage'
   | 'model.low_vram';
-type ModelGroup = 'image' | 'video';
+type ModelGroup = 'image' | 'instruction' | 'video';
 
 export interface ModelArch {
   name: string;
@@ -44,7 +44,7 @@ export const modelArchs: ModelArch[] = [
   {
     name: 'flux_kontext',
     label: 'FLUX.1-Kontext-dev',
-    group: 'image',
+    group: 'instruction',
     defaults: {
       // default updates when [selected, unselected] in the UI
       'config.process[0].model.name_or_path': ['black-forest-labs/FLUX.1-Kontext-dev', defaultNameOrPath],
@@ -306,6 +306,27 @@ export const modelArchs: ModelArch[] = [
       '3 bit with ARA': 'uint3|ostris/accuracy_recovery_adapters/qwen_image_torchao_uint3.safetensors',
     },
   },
+  {
+    name: 'qwen_image_edit',
+    label: 'Qwen-Image-Edit',
+    group: 'instruction',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['Qwen/Qwen-Image-Edit', 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].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'],
+    additionalSections: ['datasets.control_path', 'sample.ctrl_img', 'model.low_vram'],
+    accuracyRecoveryAdapters: {
+      '3 bit with ARA': 'uint3|ostris/accuracy_recovery_adapters/qwen_image_edit_torchao_uint3.safetensors',
+    },
+  },
   {
     name: 'hidream',
     label: 'HiDream',
@@ -327,7 +348,7 @@ export const modelArchs: ModelArch[] = [
   {
     name: 'hidream_e1',
     label: 'HiDream E1',
-    group: 'image',
+    group: 'instruction',
     defaults: {
       // default updates when [selected, unselected] in the UI
       'config.process[0].model.name_or_path': ['HiDream-ai/HiDream-E1-1', defaultNameOrPath],
diff --git a/version.py b/version.py
index a69dd004..e023ab46 100644
--- a/version.py
+++ b/version.py
@@ -1 +1 @@
-VERSION = "0.5.3"
\ No newline at end of file
+VERSION = "0.5.4"
\ No newline at end of file