> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.models.bagel.autoencoder

Autoencoder used by BAGEL Stage 2 image generation training.

## Module Contents

### Classes

| Name                                                                                         | Description                                                     |
| -------------------------------------------------------------------------------------------- | --------------------------------------------------------------- |
| [`AttnBlock`](#nemo_automodel-components-models-bagel-autoencoder-AttnBlock)                 | Single-head spatial attention block used in the VAE bottleneck. |
| [`AutoEncoder`](#nemo_automodel-components-models-bagel-autoencoder-AutoEncoder)             | BAGEL Stage 2 autoencoder wrapper.                              |
| [`AutoEncoderParams`](#nemo_automodel-components-models-bagel-autoencoder-AutoEncoderParams) | Architecture parameters for the BAGEL/FLUX autoencoder.         |
| [`Decoder`](#nemo_automodel-components-models-bagel-autoencoder-Decoder)                     | BAGEL/FLUX autoencoder decoder.                                 |
| [`DiagonalGaussian`](#nemo_automodel-components-models-bagel-autoencoder-DiagonalGaussian)   | Convert latent moments to a Gaussian sample or mean.            |
| [`Downsample`](#nemo_automodel-components-models-bagel-autoencoder-Downsample)               | Stride-2 downsample with explicit asymmetric padding.           |
| [`Encoder`](#nemo_automodel-components-models-bagel-autoencoder-Encoder)                     | BAGEL/FLUX autoencoder encoder.                                 |
| [`ResnetBlock`](#nemo_automodel-components-models-bagel-autoencoder-ResnetBlock)             | Residual convolution block used by the autoencoder.             |
| [`Upsample`](#nemo_automodel-components-models-bagel-autoencoder-Upsample)                   | Nearest-neighbor upsample followed by a 3x3 convolution.        |

### Functions

| Name                                                                                                           | Description                                                  |
| -------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------ |
| [`_log_load_warning`](#nemo_automodel-components-models-bagel-autoencoder-_log_load_warning)                   | -                                                            |
| [`default_autoencoder_params`](#nemo_automodel-components-models-bagel-autoencoder-default_autoencoder_params) | Return the BAGEL-7B-MoT autoencoder architecture parameters. |
| [`load_bagel_autoencoder`](#nemo_automodel-components-models-bagel-autoencoder-load_bagel_autoencoder)         | Load the BAGEL autoencoder from `ae.safetensors`.            |
| [`swish`](#nemo_automodel-components-models-bagel-autoencoder-swish)                                           | Swish activation.                                            |

### Data

[`logger`](#nemo_automodel-components-models-bagel-autoencoder-logger)

### API

```python
class nemo_automodel.components.models.bagel.autoencoder.AttnBlock(
    in_channels: int
)
```

**Bases:** `Module`

Single-head spatial attention block used in the VAE bottleneck.

```python
nemo_automodel.components.models.bagel.autoencoder.AttnBlock.attention(
    h_: torch.Tensor
) -> torch.Tensor
```

Apply scaled dot-product attention over flattened image positions.

```python
nemo_automodel.components.models.bagel.autoencoder.AttnBlock.forward(
    x: torch.Tensor
) -> torch.Tensor
```

Apply residual attention.

```python
class nemo_automodel.components.models.bagel.autoencoder.AutoEncoder(
    params: nemo_automodel.components.models.bagel.autoencoder.AutoEncoderParams
)
```

**Bases:** `Module`

BAGEL Stage 2 autoencoder wrapper.

```python
nemo_automodel.components.models.bagel.autoencoder.AutoEncoder.decode(
    z: torch.Tensor
) -> torch.Tensor
```

Decode scaled latents to image tensors.

```python
nemo_automodel.components.models.bagel.autoencoder.AutoEncoder.encode(
    x: torch.Tensor
) -> torch.Tensor
```

Encode image tensors to scaled latents.

```python
nemo_automodel.components.models.bagel.autoencoder.AutoEncoder.forward(
    x: torch.Tensor
) -> torch.Tensor
```

Encode and decode image tensors.

```python
class nemo_automodel.components.models.bagel.autoencoder.AutoEncoderParams(
    resolution: int,
    in_channels: int,
    downsample: int,
    ch: int,
    out_ch: int,
    ch_mult: list[int],
    num_res_blocks: int,
    z_channels: int,
    scale_factor: float,
    shift_factor: float
)
```

Dataclass

Architecture parameters for the BAGEL/FLUX autoencoder.

```python
class nemo_automodel.components.models.bagel.autoencoder.Decoder(
    ch: int,
    out_ch: int,
    ch_mult: list[int],
    num_res_blocks: int,
    in_channels: int,
    resolution: int,
    z_channels: int
)
```

**Bases:** `Module`

BAGEL/FLUX autoencoder decoder.

```python
nemo_automodel.components.models.bagel.autoencoder.Decoder.forward(
    z: torch.Tensor
) -> torch.Tensor
```

Decode latents to image tensors.

```python
class nemo_automodel.components.models.bagel.autoencoder.DiagonalGaussian(
    sample: bool = True,
    chunk_dim: int = 1
)
```

**Bases:** `Module`

Convert latent moments to a Gaussian sample or mean.

```python
nemo_automodel.components.models.bagel.autoencoder.DiagonalGaussian.forward(
    z: torch.Tensor
) -> torch.Tensor
```

Sample from or return the mean of a diagonal Gaussian.

```python
class nemo_automodel.components.models.bagel.autoencoder.Downsample(
    in_channels: int
)
```

**Bases:** `Module`

Stride-2 downsample with explicit asymmetric padding.

```python
nemo_automodel.components.models.bagel.autoencoder.Downsample.forward(
    x: torch.Tensor
) -> torch.Tensor
```

Downsample spatial dimensions by 2.

```python
class nemo_automodel.components.models.bagel.autoencoder.Encoder(
    resolution: int,
    in_channels: int,
    ch: int,
    ch_mult: list[int],
    num_res_blocks: int,
    z_channels: int
)
```

**Bases:** `Module`

BAGEL/FLUX autoencoder encoder.

```python
nemo_automodel.components.models.bagel.autoencoder.Encoder.forward(
    x: torch.Tensor
) -> torch.Tensor
```

Encode an image tensor to Gaussian latent moments.

```python
class nemo_automodel.components.models.bagel.autoencoder.ResnetBlock(
    in_channels: int,
    out_channels: int
)
```

**Bases:** `Module`

Residual convolution block used by the autoencoder.

```python
nemo_automodel.components.models.bagel.autoencoder.ResnetBlock.forward(
    x: torch.Tensor
) -> torch.Tensor
```

Run the residual block.

```python
class nemo_automodel.components.models.bagel.autoencoder.Upsample(
    in_channels: int
)
```

**Bases:** `Module`

Nearest-neighbor upsample followed by a 3x3 convolution.

```python
nemo_automodel.components.models.bagel.autoencoder.Upsample.forward(
    x: torch.Tensor
) -> torch.Tensor
```

Upsample spatial dimensions by 2.

```python
nemo_automodel.components.models.bagel.autoencoder._log_load_warning(
    missing: list[str],
    unexpected: list[str]
) -> None
```

```python
nemo_automodel.components.models.bagel.autoencoder.default_autoencoder_params() -> nemo_automodel.components.models.bagel.autoencoder.AutoEncoderParams
```

Return the BAGEL-7B-MoT autoencoder architecture parameters.

```python
nemo_automodel.components.models.bagel.autoencoder.load_bagel_autoencoder(
    local_path: str | None
) -> tuple[nemo_automodel.components.models.bagel.autoencoder.AutoEncoder, nemo_automodel.components.models.bagel.autoencoder.AutoEncoderParams]
```

Load the BAGEL autoencoder from `ae.safetensors`.

**Parameters:**

Local path to `ae.safetensors`. If `None`, the module is
returned with randomly initialized weights.

**Returns:** `tuple[AutoEncoder, AutoEncoderParams]`

The autoencoder module and its architecture parameters.

```python
nemo_automodel.components.models.bagel.autoencoder.swish(
    x: torch.Tensor
) -> torch.Tensor
```

Swish activation.

```python
nemo_automodel.components.models.bagel.autoencoder.logger = logging.getLogger(__name__)
```