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

patrickvonplaten · patil-suraj · web-flow · commit bf5ca036fa7f · 2022-09-19T16:00:54.000+02:00
* [Flax] Add Vae

* correct

* Apply suggestions from code review

Co-authored-by: Suraj Patil &lt;surajp815@gmail.com&gt;

* Finish

Co-authored-by: Suraj Patil &lt;surajp815@gmail.com&gt;
diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py
@@ -65,6 +65,7 @@
 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,
diff --git a/src/diffusers/modeling_flax_utils.py b/src/diffusers/modeling_flax_utils.py
@@ -294,6 +294,7 @@ def from_pretrained(
             local_files_only=local_files_only,
             use_auth_token=use_auth_token,
             revision=revision,
+            subfolder=subfolder,
             # model args
             dtype=dtype,
             **kwargs,
diff --git a/src/diffusers/models/vae_flax.py b/src/diffusers/models/vae_flax.py
