Skip to content

Commit 56b3b21

Browse files
sayakpaulsayakpaul
authored
[Refactor autoencoders] feat: introduce autoencoders module (#6129)
* feat: introduce autoencoders module * more changes for styling and copy fixing * path changes in the docs. * fix: import structure in init. * fix controlnetxs import
1 parent 9cef07d commit 56b3b21

File tree

15 files changed

+68
-61
lines changed

15 files changed

+68
-61
lines changed

‎docs/source/en/api/models/asymmetricautoencoderkl.md‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -49,12 +49,12 @@ make_image_grid([original_image, mask_image, image], rows=1, cols=3)
4949

5050
## AsymmetricAutoencoderKL
5151

52-
[[autodoc]] models.autoencoder_asym_kl.AsymmetricAutoencoderKL
52+
[[autodoc]] models.autoencoders.autoencoder_asym_kl.AsymmetricAutoencoderKL
5353

5454
## AutoencoderKLOutput
5555

56-
[[autodoc]] models.autoencoder_kl.AutoencoderKLOutput
56+
[[autodoc]] models.autoencoders.autoencoder_kl.AutoencoderKLOutput
5757

5858
## DecoderOutput
5959

60-
[[autodoc]] models.vae.DecoderOutput
60+
[[autodoc]] models.autoencoders.vae.DecoderOutput

‎docs/source/en/api/models/autoencoder_tiny.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,4 +54,4 @@ image
5454

5555
## AutoencoderTinyOutput
5656

57-
[[autodoc]] models.autoencoder_tiny.AutoencoderTinyOutput
57+
[[autodoc]] models.autoencoders.autoencoder_tiny.AutoencoderTinyOutput

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,11 +36,11 @@ model = AutoencoderKL.from_single_file(url)
3636

3737
## AutoencoderKLOutput
3838

39-
[[autodoc]] models.autoencoder_kl.AutoencoderKLOutput
39+
[[autodoc]] models.autoencoders.autoencoder_kl.AutoencoderKLOutput
4040

4141
## DecoderOutput
4242

43-
[[autodoc]] models.vae.DecoderOutput
43+
[[autodoc]] models.autoencoders.vae.DecoderOutput
4444

4545
## FlaxAutoencoderKL
4646

‎scripts/convert_consistency_decoder.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,9 @@
1212
from tqdm import tqdm
1313

1414
from diffusers import AutoencoderKL, ConsistencyDecoderVAE, DiffusionPipeline, StableDiffusionPipeline, UNet2DModel
15+
from diffusers.models.autoencoders.vae import Encoder
1516
from diffusers.models.embeddings import TimestepEmbedding
1617
from diffusers.models.unet_2d_blocks import ResnetDownsampleBlock2D, ResnetUpsampleBlock2D, UNetMidBlock2D
17-
from diffusers.models.vae import Encoder
1818

1919

2020
args = ArgumentParser()

‎src/diffusers/models/__init__.py‎

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -26,11 +26,11 @@
2626

2727
if is_torch_available():
2828
_import_structure["adapter"] = ["MultiAdapter", "T2IAdapter"]
29-
_import_structure["autoencoder_asym_kl"] = ["AsymmetricAutoencoderKL"]
30-
_import_structure["autoencoder_kl"] = ["AutoencoderKL"]
31-
_import_structure["autoencoder_kl_temporal_decoder"] = ["AutoencoderKLTemporalDecoder"]
32-
_import_structure["autoencoder_tiny"] = ["AutoencoderTiny"]
33-
_import_structure["consistency_decoder_vae"] = ["ConsistencyDecoderVAE"]
29+
_import_structure["autoencoders.autoencoder_asym_kl"] = ["AsymmetricAutoencoderKL"]
30+
_import_structure["autoencoders.autoencoder_kl"] = ["AutoencoderKL"]
31+
_import_structure["autoencoders.autoencoder_kl_temporal_decoder"] = ["AutoencoderKLTemporalDecoder"]
32+
_import_structure["autoencoders.autoencoder_tiny"] = ["AutoencoderTiny"]
33+
_import_structure["autoencoders.consistency_decoder_vae"] = ["ConsistencyDecoderVAE"]
3434
_import_structure["controlnet"] = ["ControlNetModel"]
3535
_import_structure["controlnetxs"] = ["ControlNetXSModel"]
3636
_import_structure["dual_transformer_2d"] = ["DualTransformer2DModel"]
@@ -58,11 +58,13 @@
5858
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
5959
if is_torch_available():
6060
from .adapter import MultiAdapter, T2IAdapter
61-
from .autoencoder_asym_kl import AsymmetricAutoencoderKL
62-
from .autoencoder_kl import AutoencoderKL
63-
from .autoencoder_kl_temporal_decoder import AutoencoderKLTemporalDecoder
64-
from .autoencoder_tiny import AutoencoderTiny
65-
from .consistency_decoder_vae import ConsistencyDecoderVAE
61+
from .autoencoders import (
62+
AsymmetricAutoencoderKL,
63+
AutoencoderKL,
64+
AutoencoderKLTemporalDecoder,
65+
AutoencoderTiny,
66+
ConsistencyDecoderVAE,
67+
)
6668
from .controlnet import ControlNetModel
6769
from .controlnetxs import ControlNetXSModel
6870
from .dual_transformer_2d import DualTransformer2DModel

‎src/diffusers/models/autoencoders/__init__.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
from .autoencoder_asym_kl import AsymmetricAutoencoderKL
2+
from .autoencoder_kl import AutoencoderKL
3+
from .autoencoder_kl_temporal_decoder import AutoencoderKLTemporalDecoder
4+
from .autoencoder_tiny import AutoencoderTiny
5+
from .consistency_decoder_vae import ConsistencyDecoderVAE

src/diffusers/models/autoencoder_asym_kl.py renamed to src/diffusers/models/autoencoders/autoencoder_asym_kl.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,10 @@
1616
import torch
1717
import torch.nn as nn
1818

19-
from ..configuration_utils import ConfigMixin, register_to_config
20-
from ..utils.accelerate_utils import apply_forward_hook
21-
from .modeling_outputs import AutoencoderKLOutput
22-
from .modeling_utils import ModelMixin
19+
from ...configuration_utils import ConfigMixin, register_to_config
20+
from ...utils.accelerate_utils import apply_forward_hook
21+
from ..modeling_outputs import AutoencoderKLOutput
22+
from ..modeling_utils import ModelMixin
2323
from .vae import DecoderOutput, DiagonalGaussianDistribution, Encoder, MaskConditionDecoder
2424

2525

src/diffusers/models/autoencoder_kl.py renamed to src/diffusers/models/autoencoders/autoencoder_kl.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,19 +16,19 @@
1616
import torch
1717
import torch.nn as nn
1818

19-
from ..configuration_utils import ConfigMixin, register_to_config
20-
from ..loaders import FromOriginalVAEMixin
21-
from ..utils.accelerate_utils import apply_forward_hook
22-
from .attention_processor import (
19+
from ...configuration_utils import ConfigMixin, register_to_config
20+
from ...loaders import FromOriginalVAEMixin
21+
from ...utils.accelerate_utils import apply_forward_hook
22+
from ..attention_processor import (
2323
ADDED_KV_ATTENTION_PROCESSORS,
2424
CROSS_ATTENTION_PROCESSORS,
2525
Attention,
2626
AttentionProcessor,
2727
AttnAddedKVProcessor,
2828
AttnProcessor,
2929
)
30-
from .modeling_outputs import AutoencoderKLOutput
31-
from .modeling_utils import ModelMixin
30+
from ..modeling_outputs import AutoencoderKLOutput
31+
from ..modeling_utils import ModelMixin
3232
from .vae import Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder
3333

3434

src/diffusers/models/autoencoder_kl_temporal_decoder.py renamed to src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,14 @@
1616
import torch
1717
import torch.nn as nn
1818

19-
from ..configuration_utils import ConfigMixin, register_to_config
20-
from ..loaders import FromOriginalVAEMixin
21-
from ..utils import is_torch_version
22-
from ..utils.accelerate_utils import apply_forward_hook
23-
from .attention_processor import CROSS_ATTENTION_PROCESSORS, AttentionProcessor, AttnProcessor
24-
from .modeling_outputs import AutoencoderKLOutput
25-
from .modeling_utils import ModelMixin
26-
from .unet_3d_blocks import MidBlockTemporalDecoder, UpBlockTemporalDecoder
19+
from ...configuration_utils import ConfigMixin, register_to_config
20+
from ...loaders import FromOriginalVAEMixin
21+
from ...utils import is_torch_version
22+
from ...utils.accelerate_utils import apply_forward_hook
23+
from ..attention_processor import CROSS_ATTENTION_PROCESSORS, AttentionProcessor, AttnProcessor
24+
from ..modeling_outputs import AutoencoderKLOutput
25+
from ..modeling_utils import ModelMixin
26+
from ..unet_3d_blocks import MidBlockTemporalDecoder, UpBlockTemporalDecoder
2727
from .vae import DecoderOutput, DiagonalGaussianDistribution, Encoder
2828

2929

src/diffusers/models/autoencoder_tiny.py renamed to src/diffusers/models/autoencoders/autoencoder_tiny.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,10 @@
1818

1919
import torch
2020

21-
from ..configuration_utils import ConfigMixin, register_to_config
22-
from ..utils import BaseOutput
23-
from ..utils.accelerate_utils import apply_forward_hook
24-
from .modeling_utils import ModelMixin
21+
from ...configuration_utils import ConfigMixin, register_to_config
22+
from ...utils import BaseOutput
23+
from ...utils.accelerate_utils import apply_forward_hook
24+
from ..modeling_utils import ModelMixin
2525
from .vae import DecoderOutput, DecoderTiny, EncoderTiny
2626

2727

0 commit comments

Comments
 (0)