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.