Unet2D model use ConfigMixin & ModelMixin

[mishig25]

Since diffusers models use ConfigMixin rather than Config Class, I've made:

1. Unet2D use ConfigMixin & removed the Unet2DConfig class entirely
2. Unet2D use ModelMixin, which is mostly copy of transformers.FlaxPretrainedModel but without the functionalities that are needed currently in diffusers (i.e. I've removed functionalities such as sharded checkpoints, pt-flax checkpint conversion, and other features that are not found in diffusers.ModelMixin currently)

Please let me know how this PR is looking 😊 If it is not a change we want, happy to just close it. Otherwise, happy to extend the changes to AutoencoderKL model & add necessary tests

cc: @patil-suraj @patrickvonplaten @pcuenca (I guess because of permissions can't add you as reviewers)

[patil-suraj]

Super cool Mishig! The changes look good!
This will now officially become part of diffusers, so let's maybe integrate it there :)

[mishig25]

oh I see, please let me know what should be the plan.

In diffusers repo, do we plan to have:

1. create branch jax_main_branch
2. create/merge various PRs into jax_main_branch
3. merge jax_main_branch into main

or will the process be something different?