[Chore] move vae_flax to autoencoders

docs/source/en/api/models/autoencoderkl.md

@@ -51,8 +51,8 @@ model = AutoencoderKL.from_single_file(url)
 
 ## FlaxAutoencoderKLOutput
 
-[[autodoc]] models.vae_flax.FlaxAutoencoderKLOutput
+[[autodoc]] models.autoencoders.vae_flax.FlaxAutoencoderKLOutput
 
 ## FlaxDecoderOutput
 
-[[autodoc]] models.vae_flax.FlaxDecoderOutput
+[[autodoc]] models.autoencoders.vae_flax.FlaxDecoderOutput

src/diffusers/__init__.py

@@ -422,10 +422,10 @@
 
 
 else:
+    _import_structure["models.autoencoders.vae_flax"] = ["FlaxAutoencoderKL"]
     _import_structure["models.controlnet_flax"] = ["FlaxControlNetModel"]
     _import_structure["models.modeling_flax_utils"] = ["FlaxModelMixin"]
     _import_structure["models.unets.unet_2d_condition_flax"] = ["FlaxUNet2DConditionModel"]
-    _import_structure["models.vae_flax"] = ["FlaxAutoencoderKL"]
     _import_structure["pipelines"].extend(["FlaxDiffusionPipeline"])
     _import_structure["schedulers"].extend(
         [
@@ -791,10 +791,10 @@
     except OptionalDependencyNotAvailable:
         from .utils.dummy_flax_objects import *  # noqa F403
     else:
+        from .models.autoencoders.vae_flax import FlaxAutoencoderKL
         from .models.controlnet_flax import FlaxControlNetModel
         from .models.modeling_flax_utils import FlaxModelMixin
         from .models.unets.unet_2d_condition_flax import FlaxUNet2DConditionModel
-        from .models.vae_flax import FlaxAutoencoderKL
         from .pipelines import FlaxDiffusionPipeline
         from .schedulers import (
             FlaxDDIMScheduler,

src/diffusers/models/__init__.py

@@ -59,10 +59,10 @@
     _import_structure["unets.uvit_2d"] = ["UVit2DModel"]
 
 if is_flax_available():
+    _import_structure["autoencoders.vae_flax"] = ["FlaxAutoencoderKL"]
+
     _import_structure["controlnet_flax"] = ["FlaxControlNetModel"]
     _import_structure["unets.unet_2d_condition_flax"] = ["FlaxUNet2DConditionModel"]
-    _import_structure["vae_flax"] = ["FlaxAutoencoderKL"]
-
 
 if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
     if is_torch_available():
@@ -107,9 +107,9 @@
         )
 
     if is_flax_available():
+        from .autoencoders import FlaxAutoencoderKL
         from .controlnet_flax import FlaxControlNetModel
         from .unets import FlaxUNet2DConditionModel
-        from .vae_flax import FlaxAutoencoderKL
 
 else:
     import sys

src/diffusers/models/autoencoders/__init__.py

@@ -1,6 +1,14 @@
-from .autoencoder_asym_kl import AsymmetricAutoencoderKL
-from .autoencoder_kl import AutoencoderKL
-from .autoencoder_kl_temporal_decoder import AutoencoderKLTemporalDecoder
-from .autoencoder_tiny import AutoencoderTiny
-from .consistency_decoder_vae import ConsistencyDecoderVAE
-from .vq_model import VQModel
+from ...utils import is_flax_available, is_torch_available
+
+
+if is_torch_available():
+    from .autoencoder_asym_kl import AsymmetricAutoencoderKL
+    from .autoencoder_kl import AutoencoderKL
+    from .autoencoder_kl_temporal_decoder import AutoencoderKLTemporalDecoder
+    from .autoencoder_tiny import AutoencoderTiny
+    from .consistency_decoder_vae import ConsistencyDecoderVAE
+    from .vq_model import VQModel
+
+
+if is_flax_available():
+    from .vae_flax import FlaxAutoencoderKL

src/diffusers/models/vae_flax.py → src/diffusers/models/autoencoders/vae_flax.py

@@ -24,9 +24,9 @@
 import jax.numpy as jnp
 from flax.core.frozen_dict import FrozenDict
 
-from ..configuration_utils import ConfigMixin, flax_register_to_config
-from ..utils import BaseOutput
-from .modeling_flax_utils import FlaxModelMixin
+from ...configuration_utils import ConfigMixin, flax_register_to_config
+from ...utils import BaseOutput
+from ..modeling_flax_utils import FlaxModelMixin
 
 
 @flax.struct.dataclass

src/diffusers/utils/dummy_flax_objects.py

@@ -2,7 +2,7 @@
 from ..utils import DummyObject, requires_backends
 
 
-class FlaxControlNetModel(metaclass=DummyObject):
+class FlaxAutoencoderKL(metaclass=DummyObject):
     _backends = ["flax"]
 
     def __init__(self, *args, **kwargs):
@@ -17,7 +17,7 @@ def from_pretrained(cls, *args, **kwargs):
         requires_backends(cls, ["flax"])
 
 
-class FlaxModelMixin(metaclass=DummyObject):
+class FlaxControlNetModel(metaclass=DummyObject):
     _backends = ["flax"]
 
     def __init__(self, *args, **kwargs):
@@ -32,7 +32,7 @@ def from_pretrained(cls, *args, **kwargs):
         requires_backends(cls, ["flax"])
 
 
-class FlaxUNet2DConditionModel(metaclass=DummyObject):
+class FlaxModelMixin(metaclass=DummyObject):
     _backends = ["flax"]
 
     def __init__(self, *args, **kwargs):
@@ -47,7 +47,7 @@ def from_pretrained(cls, *args, **kwargs):
         requires_backends(cls, ["flax"])
 
 
-class FlaxAutoencoderKL(metaclass=DummyObject):
+class FlaxUNet2DConditionModel(metaclass=DummyObject):
     _backends = ["flax"]
 
     def __init__(self, *args, **kwargs):