[Flax] Add Vae for Stable Diffusion (#555)

* [Flax] Add Vae

* correct

* Apply suggestions from code review
Co-authored-by: Suraj Patil <surajp815@gmail.com>

* Finish
Co-authored-by: Suraj Patil <surajp815@gmail.com>

src/diffusers/__init__.py

@@ -65,6 +65,7 @@ else:
 if is_flax_available():
     from .modeling_flax_utils import FlaxModelMixin
     from .models.unet_2d_condition_flax import FlaxUNet2DConditionModel
+    from .models.vae_flax import FlaxAutoencoderKL
     from .schedulers import (
         FlaxDDIMScheduler,
         FlaxDDPMScheduler,

src/diffusers/modeling_flax_utils.py

@@ -294,6 +294,7 @@ class FlaxModelMixin:
             local_files_only=local_files_only,
             use_auth_token=use_auth_token,
             revision=revision,
+            subfolder=subfolder,
             # model args
             dtype=dtype,
             **kwargs,