autoencoder_kl.py 24.6 KB
Newer Older
1
# Copyright 2024 The HuggingFace Team. All rights reserved.
2
3
4
5
6
7
8
9
10
11
12
13
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
14
from typing import Dict, Optional, Tuple, Union
15
16
17
18

import torch
import torch.nn as nn

19
from ...configuration_utils import ConfigMixin, register_to_config
20
from ...loaders import PeftAdapterMixin
21
from ...loaders.single_file_model import FromOriginalModelMixin
22
from ...utils import deprecate
23
24
from ...utils.accelerate_utils import apply_forward_hook
from ..attention_processor import (
25
26
    ADDED_KV_ATTENTION_PROCESSORS,
    CROSS_ATTENTION_PROCESSORS,
27
    Attention,
28
29
30
    AttentionProcessor,
    AttnAddedKVProcessor,
    AttnProcessor,
31
    FusedAttnProcessor2_0,
32
)
33
34
from ..modeling_outputs import AutoencoderKLOutput
from ..modeling_utils import ModelMixin
35
36
37
from .vae import Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder


38
class AutoencoderKL(ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapterMixin):
Steven Liu's avatar
Steven Liu committed
39
40
    r"""
    A VAE model with KL loss for encoding images into latents and decoding latent representations into images.
41

Steven Liu's avatar
Steven Liu committed
42
43
    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
    for all models (such as downloading or saving).
44
45
46
47

    Parameters:
        in_channels (int, *optional*, defaults to 3): Number of channels in the input image.
        out_channels (int,  *optional*, defaults to 3): Number of channels in the output.
Steven Liu's avatar
Steven Liu committed
48
49
50
51
52
53
        down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
            Tuple of downsample block types.
        up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
            Tuple of upsample block types.
        block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`):
            Tuple of block output channels.
54
        act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
55
        latent_channels (`int`, *optional*, defaults to 4): Number of channels in the latent space.
Steven Liu's avatar
Steven Liu committed
56
        sample_size (`int`, *optional*, defaults to `32`): Sample input size.
57
58
59
60
61
62
63
        scaling_factor (`float`, *optional*, defaults to 0.18215):
            The component-wise standard deviation of the trained latent space computed using the first batch of the
            training set. This is used to scale the latent space to have unit variance when training the diffusion
            model. The latents are scaled with the formula `z = z * scaling_factor` before being passed to the
            diffusion model. When decoding, the latents are scaled back to the original scale with the formula: `z = 1
            / scaling_factor * z`. For more details, refer to sections 4.3.2 and D.1 of the [High-Resolution Image
            Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) paper.
64
65
66
67
        force_upcast (`bool`, *optional*, default to `True`):
            If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE
            can be fine-tuned / trained to a lower range without loosing too much precision in which case
            `force_upcast` can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix
68
69
70
        mid_block_add_attention (`bool`, *optional*, default to `True`):
            If enabled, the mid_block of the Encoder and Decoder will have attention blocks. If set to false, the
            mid_block will only have resnet blocks
71
72
    """

73
    _supports_gradient_checkpointing = True
74
    _no_split_modules = ["BasicTransformerBlock", "ResnetBlock2D"]
75

76
77
78
79
80
81
82
83
84
85
86
87
88
    @register_to_config
    def __init__(
        self,
        in_channels: int = 3,
        out_channels: int = 3,
        down_block_types: Tuple[str] = ("DownEncoderBlock2D",),
        up_block_types: Tuple[str] = ("UpDecoderBlock2D",),
        block_out_channels: Tuple[int] = (64,),
        layers_per_block: int = 1,
        act_fn: str = "silu",
        latent_channels: int = 4,
        norm_num_groups: int = 32,
        sample_size: int = 32,
89
        scaling_factor: float = 0.18215,
Dhruv Nair's avatar
Dhruv Nair committed
90
        shift_factor: Optional[float] = None,
91
92
        latents_mean: Optional[Tuple[float]] = None,
        latents_std: Optional[Tuple[float]] = None,
93
        force_upcast: float = True,
Dhruv Nair's avatar
Dhruv Nair committed
94
95
        use_quant_conv: bool = True,
        use_post_quant_conv: bool = True,
96
        mid_block_add_attention: bool = True,
97
98
99
100
101
102
103
104
105
106
107
108
109
    ):
        super().__init__()

        # pass init params to Encoder
        self.encoder = Encoder(
            in_channels=in_channels,
            out_channels=latent_channels,
            down_block_types=down_block_types,
            block_out_channels=block_out_channels,
            layers_per_block=layers_per_block,
            act_fn=act_fn,
            norm_num_groups=norm_num_groups,
            double_z=True,
110
            mid_block_add_attention=mid_block_add_attention,
111
112
113
114
115
116
117
118
119
120
121
        )

        # pass init params to Decoder
        self.decoder = Decoder(
            in_channels=latent_channels,
            out_channels=out_channels,
            up_block_types=up_block_types,
            block_out_channels=block_out_channels,
            layers_per_block=layers_per_block,
            norm_num_groups=norm_num_groups,
            act_fn=act_fn,
122
            mid_block_add_attention=mid_block_add_attention,
123
124
        )

Dhruv Nair's avatar
Dhruv Nair committed
125
126
        self.quant_conv = nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1) if use_quant_conv else None
        self.post_quant_conv = nn.Conv2d(latent_channels, latent_channels, 1) if use_post_quant_conv else None
127
128
129
130
131
132
133
134
135
136
137

        self.use_slicing = False
        self.use_tiling = False

        # only relevant if vae tiling is enabled
        self.tile_sample_min_size = self.config.sample_size
        sample_size = (
            self.config.sample_size[0]
            if isinstance(self.config.sample_size, (list, tuple))
            else self.config.sample_size
        )
138
        self.tile_latent_min_size = int(sample_size / (2 ** (len(self.config.block_out_channels) - 1)))
139
140
        self.tile_overlap_factor = 0.25

141
142
143
144
    def _set_gradient_checkpointing(self, module, value=False):
        if isinstance(module, (Encoder, Decoder)):
            module.gradient_checkpointing = value

145
146
147
    def enable_tiling(self, use_tiling: bool = True):
        r"""
        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
Steven Liu's avatar
Steven Liu committed
148
149
        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
        processing larger images.
150
151
152
153
154
        """
        self.use_tiling = use_tiling

    def disable_tiling(self):
        r"""
Steven Liu's avatar
Steven Liu committed
155
156
        Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing
        decoding in one step.
157
158
159
160
161
162
163
164
165
166
167
168
        """
        self.enable_tiling(False)

    def enable_slicing(self):
        r"""
        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
        """
        self.use_slicing = True

    def disable_slicing(self):
        r"""
Steven Liu's avatar
Steven Liu committed
169
        Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing
170
171
        decoding in one step.
        """
172
173
        self.use_slicing = False

174
    @property
175
    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
176
177
178
179
180
181
182
183
184
185
    def attn_processors(self) -> Dict[str, AttentionProcessor]:
        r"""
        Returns:
            `dict` of attention processors: A dictionary containing all attention processors used in the model with
            indexed by its weight name.
        """
        # set recursively
        processors = {}

        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
186
            if hasattr(module, "get_processor"):
187
                processors[f"{name}.processor"] = module.get_processor()
188
189
190
191
192
193
194
195
196
197
198

            for sub_name, child in module.named_children():
                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)

            return processors

        for name, module in self.named_children():
            fn_recursive_add_processors(name, module, processors)

        return processors

199
    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
200
    def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
201
        r"""
Steven Liu's avatar
Steven Liu committed
202
203
        Sets the attention processor to use to compute attention.

204
        Parameters:
Steven Liu's avatar
Steven Liu committed
205
            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
206
                The instantiated processor class or a dictionary of processor classes that will be set as the processor
Steven Liu's avatar
Steven Liu committed
207
208
209
210
                for **all** `Attention` layers.

                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
                processor. This is strongly recommended when setting trainable attention processors.
211
212
213
214
215
216
217
218
219
220
221
222
223

        """
        count = len(self.attn_processors.keys())

        if isinstance(processor, dict) and len(processor) != count:
            raise ValueError(
                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
            )

        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
            if hasattr(module, "set_processor"):
                if not isinstance(processor, dict):
224
                    module.set_processor(processor)
225
                else:
226
                    module.set_processor(processor.pop(f"{name}.processor"))
227
228
229
230
231
232
233

            for sub_name, child in module.named_children():
                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)

        for name, module in self.named_children():
            fn_recursive_attn_processor(name, module, processor)

234
    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor
235
236
237
238
    def set_default_attn_processor(self):
        """
        Disables custom attention processors and sets the default attention implementation.
        """
239
240
241
242
243
244
245
246
247
        if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
            processor = AttnAddedKVProcessor()
        elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
            processor = AttnProcessor()
        else:
            raise ValueError(
                f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
            )

248
        self.set_attn_processor(processor)
249

250
251
252
253
254
255
256
257
258
259
260
261
    def _encode(self, x: torch.Tensor) -> torch.Tensor:
        batch_size, num_channels, height, width = x.shape

        if self.use_tiling and (width > self.tile_sample_min_size or height > self.tile_sample_min_size):
            return self._tiled_encode(x)

        enc = self.encoder(x)
        if self.quant_conv is not None:
            enc = self.quant_conv(enc)

        return enc

262
    @apply_forward_hook
263
    def encode(
264
        self, x: torch.Tensor, return_dict: bool = True
265
266
267
268
269
    ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
        """
        Encode a batch of images into latents.

        Args:
270
            x (`torch.Tensor`): Input batch of images.
271
            return_dict (`bool`, *optional*, defaults to `True`):
Dhruv Nair's avatar
Dhruv Nair committed
272
                Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
273
274
275

        Returns:
                The latent representations of the encoded images. If `return_dict` is True, a
Dhruv Nair's avatar
Dhruv Nair committed
276
                [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
277
        """
Patrick von Platen's avatar
Patrick von Platen committed
278
        if self.use_slicing and x.shape[0] > 1:
279
            encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)]
Patrick von Platen's avatar
Patrick von Platen committed
280
281
            h = torch.cat(encoded_slices)
        else:
282
            h = self._encode(x)
Dhruv Nair's avatar
Dhruv Nair committed
283

284
        posterior = DiagonalGaussianDistribution(h)
285
286
287
288
289
290

        if not return_dict:
            return (posterior,)

        return AutoencoderKLOutput(latent_dist=posterior)

291
    def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
292
293
294
        if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):
            return self.tiled_decode(z, return_dict=return_dict)

Dhruv Nair's avatar
Dhruv Nair committed
295
296
297
        if self.post_quant_conv is not None:
            z = self.post_quant_conv(z)

298
299
300
301
302
303
304
        dec = self.decoder(z)

        if not return_dict:
            return (dec,)

        return DecoderOutput(sample=dec)

305
    @apply_forward_hook
Dhruv Nair's avatar
Dhruv Nair committed
306
307
308
    def decode(
        self, z: torch.FloatTensor, return_dict: bool = True, generator=None
    ) -> Union[DecoderOutput, torch.FloatTensor]:
309
310
311
312
        """
        Decode a batch of images.

        Args:
313
            z (`torch.Tensor`): Input batch of latent vectors.
314
315
316
317
318
319
320
321
322
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.

        Returns:
            [`~models.vae.DecoderOutput`] or `tuple`:
                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
                returned.

        """
323
324
325
326
        if self.use_slicing and z.shape[0] > 1:
            decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
            decoded = torch.cat(decoded_slices)
        else:
Dhruv Nair's avatar
Dhruv Nair committed
327
            decoded = self._decode(z).sample
328
329
330
331
332
333

        if not return_dict:
            return (decoded,)

        return DecoderOutput(sample=decoded)

334
    def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
335
336
        blend_extent = min(a.shape[2], b.shape[2], blend_extent)
        for y in range(blend_extent):
337
338
339
            b[:, :, y, :] = a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
        return b

340
    def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
341
342
        blend_extent = min(a.shape[3], b.shape[3], blend_extent)
        for x in range(blend_extent):
343
344
345
            b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
        return b

346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
    def _tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
        r"""Encode a batch of images using a tiled encoder.

        When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several
        steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is
        different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the
        tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the
        output, but they should be much less noticeable.

        Args:
            x (`torch.Tensor`): Input batch of images.

        Returns:
            `torch.Tensor`:
                The latent representation of the encoded videos.
        """

        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
        row_limit = self.tile_latent_min_size - blend_extent

        # Split the image into 512x512 tiles and encode them separately.
        rows = []
        for i in range(0, x.shape[2], overlap_size):
            row = []
            for j in range(0, x.shape[3], overlap_size):
                tile = x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
                tile = self.encoder(tile)
                if self.config.use_quant_conv:
                    tile = self.quant_conv(tile)
                row.append(tile)
            rows.append(row)
        result_rows = []
        for i, row in enumerate(rows):
            result_row = []
            for j, tile in enumerate(row):
                # blend the above tile and the left tile
                # to the current tile and add the current tile to the result row
                if i > 0:
                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
                if j > 0:
                    tile = self.blend_h(row[j - 1], tile, blend_extent)
                result_row.append(tile[:, :, :row_limit, :row_limit])
            result_rows.append(torch.cat(result_row, dim=3))

        enc = torch.cat(result_rows, dim=2)
        return enc

394
    def tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> AutoencoderKLOutput:
395
        r"""Encode a batch of images using a tiled encoder.
396

397
        When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several
Steven Liu's avatar
Steven Liu committed
398
399
        steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is
        different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the
400
        tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the
Steven Liu's avatar
Steven Liu committed
401
402
403
        output, but they should be much less noticeable.

        Args:
404
            x (`torch.Tensor`): Input batch of images.
Steven Liu's avatar
Steven Liu committed
405
            return_dict (`bool`, *optional*, defaults to `True`):
Dhruv Nair's avatar
Dhruv Nair committed
406
                Whether or not to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
Steven Liu's avatar
Steven Liu committed
407
408

        Returns:
Dhruv Nair's avatar
Dhruv Nair committed
409
410
411
            [`~models.autoencoder_kl.AutoencoderKLOutput`] or `tuple`:
                If return_dict is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain
                `tuple` is returned.
412
        """
413
414
415
416
417
418
419
        deprecation_message = (
            "The tiled_encode implementation supporting the `return_dict` parameter is deprecated. In the future, the "
            "implementation of this method will be replaced with that of `_tiled_encode` and you will no longer be able "
            "to pass `return_dict`. You will also have to create a `DiagonalGaussianDistribution()` from the returned value."
        )
        deprecate("tiled_encode", "1.0.0", deprecation_message, standard_warn=False)

420
421
422
423
424
425
426
427
428
429
430
        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
        row_limit = self.tile_latent_min_size - blend_extent

        # Split the image into 512x512 tiles and encode them separately.
        rows = []
        for i in range(0, x.shape[2], overlap_size):
            row = []
            for j in range(0, x.shape[3], overlap_size):
                tile = x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
                tile = self.encoder(tile)
431
432
                if self.config.use_quant_conv:
                    tile = self.quant_conv(tile)
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
                row.append(tile)
            rows.append(row)
        result_rows = []
        for i, row in enumerate(rows):
            result_row = []
            for j, tile in enumerate(row):
                # blend the above tile and the left tile
                # to the current tile and add the current tile to the result row
                if i > 0:
                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
                if j > 0:
                    tile = self.blend_h(row[j - 1], tile, blend_extent)
                result_row.append(tile[:, :, :row_limit, :row_limit])
            result_rows.append(torch.cat(result_row, dim=3))

        moments = torch.cat(result_rows, dim=2)
        posterior = DiagonalGaussianDistribution(moments)

        if not return_dict:
            return (posterior,)

        return AutoencoderKLOutput(latent_dist=posterior)

456
    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
Steven Liu's avatar
Steven Liu committed
457
458
        r"""
        Decode a batch of images using a tiled decoder.
459

460
        Args:
461
            z (`torch.Tensor`): Input batch of latent vectors.
Steven Liu's avatar
Steven Liu committed
462
463
464
465
466
467
468
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.

        Returns:
            [`~models.vae.DecoderOutput`] or `tuple`:
                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
                returned.
469
470
471
472
473
474
475
476
477
478
479
480
        """
        overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
        blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
        row_limit = self.tile_sample_min_size - blend_extent

        # Split z into overlapping 64x64 tiles and decode them separately.
        # The tiles have an overlap to avoid seams between tiles.
        rows = []
        for i in range(0, z.shape[2], overlap_size):
            row = []
            for j in range(0, z.shape[3], overlap_size):
                tile = z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
481
482
                if self.config.use_post_quant_conv:
                    tile = self.post_quant_conv(tile)
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
                decoded = self.decoder(tile)
                row.append(decoded)
            rows.append(row)
        result_rows = []
        for i, row in enumerate(rows):
            result_row = []
            for j, tile in enumerate(row):
                # blend the above tile and the left tile
                # to the current tile and add the current tile to the result row
                if i > 0:
                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
                if j > 0:
                    tile = self.blend_h(row[j - 1], tile, blend_extent)
                result_row.append(tile[:, :, :row_limit, :row_limit])
            result_rows.append(torch.cat(result_row, dim=3))

        dec = torch.cat(result_rows, dim=2)
        if not return_dict:
            return (dec,)

        return DecoderOutput(sample=dec)

505
506
    def forward(
        self,
507
        sample: torch.Tensor,
508
509
510
        sample_posterior: bool = False,
        return_dict: bool = True,
        generator: Optional[torch.Generator] = None,
511
    ) -> Union[DecoderOutput, torch.Tensor]:
512
513
        r"""
        Args:
514
            sample (`torch.Tensor`): Input sample.
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
            sample_posterior (`bool`, *optional*, defaults to `False`):
                Whether to sample from the posterior.
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
        """
        x = sample
        posterior = self.encode(x).latent_dist
        if sample_posterior:
            z = posterior.sample(generator=generator)
        else:
            z = posterior.mode()
        dec = self.decode(z).sample

        if not return_dict:
            return (dec,)

        return DecoderOutput(sample=dec)
532

533
    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections
534
535
    def fuse_qkv_projections(self):
        """
536
537
        Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
        are fused. For cross-attention modules, key and value projection matrices are fused.
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556

        <Tip warning={true}>

        This API is 🧪 experimental.

        </Tip>
        """
        self.original_attn_processors = None

        for _, attn_processor in self.attn_processors.items():
            if "Added" in str(attn_processor.__class__.__name__):
                raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")

        self.original_attn_processors = self.attn_processors

        for module in self.modules():
            if isinstance(module, Attention):
                module.fuse_projections(fuse=True)

557
558
        self.set_attn_processor(FusedAttnProcessor2_0())

559
    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections
560
561
562
563
564
565
566
567
568
569
570
571
    def unfuse_qkv_projections(self):
        """Disables the fused QKV projection if enabled.

        <Tip warning={true}>

        This API is 🧪 experimental.

        </Tip>

        """
        if self.original_attn_processors is not None:
            self.set_attn_processor(self.original_attn_processors)