Support training SD V2 with Flax

[yasyf]

Mostly involves supporting a v_prediction scheduler.

The implementation in #1777 doesn't take into account a recent refactor of scheduling_utils_flax, so this should be used instead.

[yasyf]

cc @patrickvonplaten @yiyixuxu @patil-suraj

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[bonlime]

Hey, sorry to ask this here, but I don't know other place for it.

Is v_prediction scheduler only used for SD-2-768 version?

I'm confused that on hugging-face config the 2-1-base(512) uses epsilon link and only 768 uses v_prediction link

[patil-suraj]

Hey @bonlime that's correct, only the 768 and upscaler models use v-prediction.

[yasyf]

@patil-suraj done!

[pcuenca]

Fixes #1798