Add hidream to the ui

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

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

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',

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,

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 {

version.py

@@ -1 +1 @@
-VERSION = "0.2.5"
+VERSION = "0.2.6"