Flax dreambooth/training script currently does not support v_prediction and thus doesn't work on SD v2

[Lime-Cakes]

Is your feature request related to a problem? Please describe.
Flax version of the training script under example did not implement different loss target for v_prediction.

Describe the solution you'd like
Implement v_prediction

Describe alternatives you've considered
I could probably implement it based on pytorch code, but I haven't had time to read and fully understand the paper, so might mess up. Just filing this here so people know flax version is different from pytorch version and will break when training on v_prediction models.

[patrickvonplaten]

Hey @Lime-Cakes,

Yes that's very true! It would be great if someone finds the time to adapt it, I'm currently very much under water with the PyTorch-side of things 😅

[Lime-Cakes]

Hey @Lime-Cakes,

Yes that's very true! It would be great if someone finds the time to adapt it, I'm currently very much under water with the PyTorch-side of things sweat_smile

#1783 seemed to fix this issue. It looked ready and had been reviewed, but wasn't yet merged.

[patrickvonplaten]

Cool, think we can close this one now no? :-)

[Lime-Cakes]

Yep!