From 79c87701e7fc6f47f1ea8adc75834284b007d3ab Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Wed, 16 Apr 2025 13:45:21 -0600
Subject: [PATCH] Add hidream to the ui

---
 config/examples/train_lora_hidream_48.yaml |  3 ++-
 ui/src/app/jobs/new/SimpleJob.tsx          |  2 +-
 ui/src/app/jobs/new/jobConfig.ts           |  9 ++++++---
 ui/src/app/jobs/new/options.ts             | 15 +++++++++++++++
 ui/src/types.ts                            |  3 +++
 version.py                                 |  2 +-
 6 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/config/examples/train_lora_hidream_48.yaml b/config/examples/train_lora_hidream_48.yaml
index ea9c7dff..f915f4c7 100644
--- a/config/examples/train_lora_hidream_48.yaml
+++ b/config/examples/train_lora_hidream_48.yaml
@@ -29,7 +29,8 @@ config:
           # it is probably best to ignore the mixture of experts since only 2 are active each block. It works activating it, but I wouldnt.
           # proper training of it is not fully implemented
           ignore_if_contains:
-            - "ff_i"
+            - "ff_i.experts"
+            - "ff_i.gate"
       save:
         dtype: bfloat16 # precision to save
         save_every: 250 # save every this many steps
diff --git a/ui/src/app/jobs/new/SimpleJob.tsx b/ui/src/app/jobs/new/SimpleJob.tsx
index 75bab27d..84d98330 100644
--- a/ui/src/app/jobs/new/SimpleJob.tsx
+++ b/ui/src/app/jobs/new/SimpleJob.tsx
@@ -283,7 +283,7 @@ export default function SimpleJob({
                   options={[
                     { value: 'sigmoid', label: 'Sigmoid' },
                     { value: 'linear', label: 'Linear' },
-                    { value: 'flux_shift', label: 'Flux Shift' },
+                    { value: 'shift', label: 'Shift' },
                   ]}
                 />
                 <SelectInput
diff --git a/ui/src/app/jobs/new/jobConfig.ts b/ui/src/app/jobs/new/jobConfig.ts
index 5f090479..cf79092f 100644
--- a/ui/src/app/jobs/new/jobConfig.ts
+++ b/ui/src/app/jobs/new/jobConfig.ts
@@ -16,7 +16,7 @@ export const defaultDatasetConfig: DatasetConfig = {
 export const defaultJobConfig: JobConfig = {
   job: 'extension',
   config: {
-    name: 'my_first_flex_lora_v1',
+    name: 'my_first_lora_v1',
     process: [
       {
         type: 'ui_trainer',
@@ -31,6 +31,9 @@ export const defaultJobConfig: JobConfig = {
           linear_alpha: 32,
           lokr_full_rank: true,
           lokr_factor: -1,
+          network_kwargs: {
+            ignore_if_contains: [],
+          },
         },
         save: {
           dtype: 'bf16',
@@ -43,7 +46,7 @@ export const defaultJobConfig: JobConfig = {
         train: {
           batch_size: 1,
           bypass_guidance_embedding: true,
-          steps: 2000,
+          steps: 3000,
           gradient_accumulation: 1,
           train_unet: true,
           train_text_encoder: false,
@@ -58,7 +61,7 @@ export const defaultJobConfig: JobConfig = {
           unload_text_encoder: false,
           lr: 0.0001,
           ema_config: {
-            use_ema: true,
+            use_ema: false,
             ema_decay: 0.99,
           },
           dtype: 'bf16',
diff --git a/ui/src/app/jobs/new/options.ts b/ui/src/app/jobs/new/options.ts
index 17c05b79..89fcf764 100644
--- a/ui/src/app/jobs/new/options.ts
+++ b/ui/src/app/jobs/new/options.ts
@@ -12,6 +12,7 @@ export const modelArchs = [
   { name: 'flux', label: 'Flux.1' },
   { name: 'wan21', label: 'Wan 2.1' },
   { name: 'lumina2', label: 'Lumina2' },
+  { name: 'hidream', label: 'HiDream' },
 ];
 
 export const isVideoModelFromArch = (arch: string) => {
@@ -83,6 +84,20 @@ export const options = {
         'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
       },
     },
+    {
+      name_or_path: 'HiDream-ai/HiDream-I1-Full',
+      defaults: {
+        // default updates when [selected, unselected] in the UI
+        'config.process[0].model.quantize': [true, false],
+        'config.process[0].model.quantize_te': [true, false],
+        'config.process[0].model.arch': ['hidream', defaultModelArch],
+        'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+        'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+        'config.process[0].train.lr': [0.0002, 0.0001],
+        'config.process[0].train.timestep_type': ['shift', 'sigmoid'],
+        'config.process[0].network.network_kwargs.ignore_if_contains': [['ff_i.experts', 'ff_i.gate'], []],
+      },
+    },
     {
       name_or_path: 'ostris/objective-reality',
       dev_only: true,
diff --git a/ui/src/types.ts b/ui/src/types.ts
index 9841553e..fdac4d46 100644
--- a/ui/src/types.ts
+++ b/ui/src/types.ts
@@ -55,6 +55,9 @@ export interface NetworkConfig {
   linear_alpha: number;
   lokr_full_rank: boolean;
   lokr_factor: number;
+  network_kwargs: {
+    ignore_if_contains: string[];
+  }
 }
 
 export interface SaveConfig {
diff --git a/version.py b/version.py
index bccbcfdf..88364324 100644
--- a/version.py
+++ b/version.py
@@ -1 +1 @@
-VERSION = "0.2.5"
\ No newline at end of file
+VERSION = "0.2.6"
\ No newline at end of file