Explicit support of masked loss and schedulefree optimizers

[StrangeTcy]

ETA: this is my first massive involvement with scripts using diffusers, so I might not be getting some concepts for now, but I'm trying to learn as I go.

I'm trying to extend a script from the advanced_diffusion_training folder that deals with finetuning a dreambooth lora for flux (https://github.com/huggingface/diffusers/blob/main/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py),
but I'm trying to:

1. add support for schedulefree optimizers (primarily AdamWScheduleFree)
2. add a way to use masked loss (based on the mask images or alpha channel info). (possibly related to Full support for Flux attention masking #10194)

I'm basing both additions on the way it's handled in sd-scripts by kohya-ss (https://github.com/kohya-ss/sd-scripts/blob/sd3/flux_train_network.py is the main source of inspiration), but

1. I'm not sure I'm adding schedulefree optimizer train / eval switching in all the right places (before actual training, before sampling images, before saving checpoints, &c)
2. the masked loss part has me stumped; I think that we can use the way the dataset is constructed (DreamBoothDataset has no out-of-the-box support for alpha_mask, but DreamBoothSubset from sd-scripts does), but maybe I’m missing something

My current attempts live here: https://gist.github.com/StrangeTcy/dc15b5880dd0d0d92639fe7aba595d54

Any pointers would be welcome.

[bghira]

masked loss is not related to attn masking, and is honestly not terribly useful. i would recommend leaving it out of the training examples. i've implemented it in simpletuner and it has dubious value.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.