github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
* 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
This commit is contained in:
Jaret Burkett (Ostris)
2026-05-12 11:15:16 -06:00
committed by GitHub
parent f38de2a2fe
commit fc4d6ebf39
14 changed files with 3633 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
ui/src/app/jobs/new/SimpleJob.tsx
View File

@@ -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 => (

35
ui/src/app/jobs/new/options.ts
View File

@@ -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' });