Add support for fine-tuning Hidream O1 (#831)

* Initial support for hidream. Lora keys likely need work

* Fix saving for hidream-o1

* Remove dependence on flash attention for hidream o1

* Fix gradient checkpointing for hidream o1

* A lot of fixes for hidream. Handle loading and saving as comfy model.

* Omit layers not used in comfy. Fix issue with lora loading keys in comfy

* Version bumpo

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

@@ -380,20 +380,24 @@ export default function SimpleJob({
                 }}
                 options={transformerQuantizationOptions}
               />
-              <SelectInput
-                label="Text Encoder"
-                value={jobConfig.config.process[0].model.quantize_te ? jobConfig.config.process[0].model.qtype_te : ''}
-                onChange={value => {
-                  if (value === '') {
-                    setJobConfig(false, 'config.process[0].model.quantize_te');
-                    value = defaultQtype;
-                  } else {
-                    setJobConfig(true, 'config.process[0].model.quantize_te');
+              {!disableSections.includes('model.quantize_te') && (
+                <SelectInput
+                  label="Text Encoder"
+                  value={
+                    jobConfig.config.process[0].model.quantize_te ? jobConfig.config.process[0].model.qtype_te : ''
                   }
-                  setJobConfig(value, 'config.process[0].model.qtype_te');
-                }}
-                options={quantizationOptions}
-              />
+                  onChange={value => {
+                    if (value === '') {
+                      setJobConfig(false, 'config.process[0].model.quantize_te');
+                      value = defaultQtype;
+                    } else {
+                      setJobConfig(true, 'config.process[0].model.quantize_te');
+                    }
+                    setJobConfig(value, 'config.process[0].model.qtype_te');
+                  }}
+                  options={quantizationOptions}
+                />
+              )}
             </Card>
           )}
           {modelArch?.additionalSections?.includes('model.multistage') && (
@@ -1068,7 +1072,7 @@ export default function SimpleJob({
                           <div className="grid grid-cols-2 gap-2">
                             {[
                               [256, 512, 768, 1024],
-                              [1280, 1328, 1536],
+                              [1280, 1328, 1536, 2048],
                             ].map(resGroup => (
                               <div key={resGroup[0]} className="space-y-2">
                                 {resGroup.map(res => (

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

@@ -6,6 +6,7 @@ type Control = 'depth' | 'line' | 'pose' | 'inpaint';
 
 type DisableableSections =
   | 'model.quantize'
+  | 'model.quantize_te'
   | 'train.timestep_type'
   | 'network.conv'
   | 'trigger_word'
@@ -922,6 +923,40 @@ export const modelArchs: ModelArch[] = [
     disableSections: ['network.conv'],
     additionalSections: ['model.low_vram'],
   },
+  {
+    name: 'hidream_o1',
+    label: 'HiDream-O1',
+    group: 'image',
+    defaults: {
+      'config.process[0].model.name_or_path': ['HiDream-ai/HiDream-O1-Image', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [false, false],
+      'config.process[0].train.timestep_type': ['linear', 'sigmoid'],
+      'config.process[0].network.conv': [undefined, 16],
+      'config.process[0].network.conv_alpha': [undefined, 16],
+      'config.process[0].train.max_loss': [1.0, undefined],
+      'config.process[0].network.network_kwargs.ignore_if_contains': [['lm_head','patch_embed', 'visual'], []],
+      'config.process[0].network.transformer_only': [false, undefined],
+      'config.process[0].sample.width': [2048, 1024],
+      'config.process[0].sample.height': [2048, 1024],
+      'config.process[0].model.model_kwargs': [
+        {
+          noise_scale_inference: 8.0,
+          noise_scale: 8.0,
+        },
+        {},
+      ],
+    },
+    disableSections: [
+      'network.conv',
+      'model.quantize_te',
+      'train.unload_text_encoder',
+    ],
+    additionalSections: [
+      'model.low_vram',
+      'model.layer_offloading',
+    ],
+  },
 ].sort((a, b) => {
   // Sort by label, case-insensitive
   return a.label.localeCompare(b.label, undefined, { sensitivity: 'base' });

ui/src/types.ts

@@ -71,6 +71,7 @@ export interface NetworkConfig {
   network_kwargs: {
     ignore_if_contains: string[];
   };
+  transformer_only?: boolean;
 }
 
 export interface SaveConfig {
@@ -149,6 +150,7 @@ export interface TrainConfig {
   do_differential_guidance?: boolean;
   differential_guidance_scale?: number;
   audio_loss_multiplier?: number;
+  max_loss?: number | null;
 }
 
 export interface QuantizeKwargsConfig {