Work in progress. Relies on my fork of trl https://github.com/andjoer/trl.git (install from my repo). Training works on a H200 GPU.
example usage:
python alignprop_fluxfill.py --mixed_precision bf16 --log_with wandb --log_image_freq 1 --save_freq 10 --num_epochs 200