Repository: Project-MONAI/MONAI
File: monai/apps/generation/maisi/networks/autoencoderkl_maisi.py (Python)

Code before:
        _empty_cuda_cache(self.save_mem)

        if self.in_channels != self.out_channels:
            x = self.nin_shortcut(x)
            _empty_cuda_cache(self.save_mem)

        out = x + h
        out_tensor: torch.Tensor = convert_to_tensor(out)
        return out_tensor


class MaisiEncoder(nn.Module):
    """
    Convolutional cascade that downsamples the image into a spatial latent space.

    Args:
        spatial_dims: Number of spatial dimensions (1D, 2D, 3D).
        in_channels: Number of input channels.
        num_channels: Sequence of block output channels.
        out_channels: Number of channels in the bottom layer (latent space) of the autoencoder.
        num_res_blocks: Number of residual blocks (see ResBlock) per level.
        norm_num_groups: Number of groups for the group norm layers.
        norm_eps: Epsilon for the normalization.
        attention_levels: Indicate which level from num_channels contain an attention block.
        with_nonlocal_attn: If True, use non-local attention block.
        include_fc: whether to include the final linear layer in the attention block. Default to True.
        use_combined_linear: whether to use a single linear layer for qkv projection in the attention block, default to False.
        use_flash_attention: If True, use flash attention for a memory efficient attention mechanism.
        num_splits: Number of splits for the input tensor.
        dim_split: Dimension of splitting for the input tensor.
        norm_float16: If True, convert output of MaisiGroupNorm3D to float16 format, default to `False`.
        print_info: Whether to print information, default to `False`.
        save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`.
    """

    def __init__(
        self,
        spatial_dims: int,
        in_channels: int,
        num_channels: Sequence[int],
        out_channels: int,
        num_res_blocks: Sequence[int],
        norm_num_groups: int,
        norm_eps: float,
        attention_levels: Sequence[bool],
        num_splits: int,
        dim_split: int,
        norm_float16: bool = False,
        print_info: bool = False,
        save_mem: bool = True,
        with_nonlocal_attn: bool = True,

Review comment (suggestion):
Please also update the docstring for num_res_blocks.
`num_res_blocks: Number of residual blocks (see ResBlock) per level.`
to
`num_res_blocks: Number of residual blocks (see AEKLResBlock) per level.`

Code after:
        _empty_cuda_cache(self.save_mem)

        if self.in_channels != self.out_channels:
            x = self.nin_shortcut(x)
            _empty_cuda_cache(self.save_mem)

        out = x + h
        out_tensor: torch.Tensor = convert_to_tensor(out)
        return out_tensor


class MaisiEncoder(nn.Module):
    """
    Convolutional cascade that downsamples the image into a spatial latent space.

    Args:
        spatial_dims: Number of spatial dimensions (1D, 2D, 3D).
        in_channels: Number of input channels.
        num_channels: Sequence of block output channels.
        out_channels: Number of channels in the bottom layer (latent space) of the autoencoder.
        num_res_blocks: Number of residual blocks (see AEKLResBlock) per level.
        norm_num_groups: Number of groups for the group norm layers.
        norm_eps: Epsilon for the normalization.
        attention_levels: Indicate which level from num_channels contain an attention block.
        with_nonlocal_attn: If True, use non-local attention block.
        include_fc: whether to include the final linear layer in the attention block. Default to False.
        use_combined_linear: whether to use a single linear layer for qkv projection in the attention block, default to False.
        use_flash_attention: If True, use flash attention for a memory efficient attention mechanism.
        num_splits: Number of splits for the input tensor.
        dim_split: Dimension of splitting for the input tensor.
        norm_float16: If True, convert output of MaisiGroupNorm3D to float16 format, default to `False`.
        print_info: Whether to print information, default to `False`.
        save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`.
    """

    def __init__(
        self,
        spatial_dims: int,
        in_channels: int,
        num_channels: Sequence[int],
        out_channels: int,
        num_res_blocks: Sequence[int],
        norm_num_groups: int,
        norm_eps: float,
        attention_levels: Sequence[bool],
        num_splits: int,
        dim_split: int,
        norm_float16: bool = False,
        print_info: bool = False,
        save_mem: bool = True,
        with_nonlocal_attn: bool = True,