@Ando233 here is something to get you started or as a reference:

from dataclasses import dataclass
from math import sqrt
from typing import Optional, Tuple, Union

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import PeftAdapterMixin
from ...loaders.single_file_model import FromOriginalModelMixin
from ...utils import BaseOutput
from ...utils.accelerate_utils import apply_forward_hook
from ..attention import FeedForward
from ..attention_processor import Attention
from ..modeling_utils import ModelMixin


@dataclass
class AutoencoderRAEOuput(BaseOutput):
    """
    Output of AutoencoderRAE encoding method.

    Args:
        latent (`torch.Tensor`):
            Encoded outputs of the encoder (frozen representation encoder).
            Shape: (batch_size, hidden_size, latent_height, latent_width)
    """

    latent: torch.Tensor


class RAEDecoderOutput(BaseOutput):
    """
    Output of RAEDecoder.

    Args:
        sample (`torch.Tensor`):
            Decoded output from decoder. Shape: (batch_size, num_channels, image_height, image_width)
    """

    sample: torch.Tensor


class ViTMAEDecoderConfig:
    """Configuration for ViTMAEDecoder.

    Based on transformers.ViTMAEConfig but adapted for RAE's decoder.
    """

    def __init__(
        self,
        hidden_size: int = 768,
        num_hidden_layers: int = 12,
        num_attention_heads: int = 12,
        intermediate_size: int = 3072,
        hidden_act: str = "gelu",
        hidden_dropout_prob: float = 0.0,
        attention_probs_dropout_prob: float = 0.0,
        initializer_range: float = 0.02,
        layer_norm_eps: float = 1e-12,
        image_size: int = 224,
        patch_size: int = 16,
        num_channels: int = 3,
        qkv_bias: bool = True,
        decoder_hidden_size: int = 512,
        decoder_num_attention_heads: int = 16,
        decoder_num_hidden_layers: int = 8,
        decoder_intermediate_size: int = 2048,
        mask_ratio: float = 0.0,
        norm_pix_loss: bool = False,
        **kwargs,
    ):
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.intermediate_size = intermediate_size
        self.hidden_act = hidden_act
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.initializer_range = initializer_range
        self.layer_norm_eps = layer_norm_eps
        self.image_size = image_size
        self.patch_size = patch_size
        self.num_channels = num_channels
        self.qkv_bias = qkv_bias
        self.decoder_hidden_size = decoder_hidden_size
        self.decoder_num_attention_heads = decoder_num_attention_heads
        self.decoder_num_hidden_layers = decoder_num_hidden_layers
        self.decoder_intermediate_size = decoder_intermediate_size
        self.mask_ratio = mask_ratio
        self.norm_pix_loss = norm_pix_loss

    def to_dict(self) -> dict:
        """Convert config to dictionary for serialization."""
        return {
            "hidden_size": self.hidden_size,
            "num_hidden_layers": self.num_hidden_layers,
            "num_attention_heads": self.num_attention_heads,
            "intermediate_size": self.intermediate_size,
            "hidden_act": self.hidden_act,
            "hidden_dropout_prob": self.hidden_dropout_prob,
            "attention_probs_dropout_prob": self.attention_probs_dropout_prob,
            "initializer_range": self.initializer_range,
            "layer_norm_eps": self.layer_norm_eps,
            "image_size": self.image_size,
            "patch_size": self.patch_size,
            "num_channels": self.num_channels,
            "qkv_bias": self.qkv_bias,
            "decoder_hidden_size": self.decoder_hidden_size,
            "decoder_num_attention_heads": self.decoder_num_attention_heads,
            "decoder_num_hidden_layers": self.decoder_num_hidden_layers,
            "decoder_intermediate_size": self.decoder_intermediate_size,
            "mask_ratio": self.mask_ratio,
            "norm_pix_loss": self.norm_pix_loss,
        }


class ViTMAEDecoderLayer(nn.Module):
    """Single ViTMAE Transformer layer using diffusers Attention and FeedForward.

    Based on transformers ViTMAELayer with pre-norm attention.
    """

    def __init__(self, config: ViTMAEDecoderConfig):
        super().__init__()
        self.layernorm_before = nn.LayerNorm(config.decoder_hidden_size, eps=config.layer_norm_eps)
        self.layernorm_after = nn.LayerNorm(config.decoder_hidden_size, eps=config.layer_norm_eps)

        # Use diffusers Attention - self attention only (no cross attention)
        self.attn = Attention(
            query_dim=config.decoder_hidden_size,
            heads=config.decoder_num_attention_heads,
            dim_head=config.decoder_hidden_size // config.decoder_num_attention_heads,
            dropout=config.attention_probs_dropout_prob,
            bias=config.qkv_bias,
        )

        # Use diffusers FeedForward
        self.ff = FeedForward(
            dim=config.decoder_hidden_size,
            inner_dim=config.decoder_intermediate_size,
            dropout=config.hidden_dropout_prob,
            activation_fn="gelu",
            bias=True,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        output_attentions: bool = False,
    ) -> Tuple[torch.Tensor, ...]:
        # Self attention with pre-norm
        residual = hidden_states
        hidden_states = self.layernorm_before(hidden_states)
        hidden_states = self.attn(hidden_states)
        hidden_states = residual + hidden_states

        # Feed forward with pre-norm
        residual = hidden_states
        hidden_states = self.layernorm_after(hidden_states)
        hidden_states = self.ff(hidden_states)
        hidden_states = residual + hidden_states

        return (hidden_states,)


def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int, add_cls_token: bool = False) -> torch.Tensor:
    """
    Create 2D sin/cos positional embeddings.

    Args:
        embed_dim: Embedding dimension
        grid_size: Grid height and width
        add_cls_token: Whether to add CLS token position

    Returns:
        Positional embeddings tensor
    """
    grid_h = np.arange(grid_size, dtype=np.float32)
    grid_w = np.arange(grid_size, dtype=np.float32)
    grid = np.meshgrid(grid_w, grid_h)
    grid = np.stack(grid, axis=0)
    grid = grid.reshape([2, 1, grid_size, grid_size])

    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
    if add_cls_token:
        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
    return torch.from_numpy(pos_embed).float()


def get_2d_sincos_pos_embed_from_grid(embed_dim: int, grid) -> np.ndarray:
    """Get 2D sin/cos pos embed from grid."""
    if embed_dim % 2 != 0:
        raise ValueError("embed_dim must be even")

    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
    emb = np.concatenate([emb_h, emb_w], axis=1)
    return emb


def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos) -> np.ndarray:
    """Get 1D sin/cos pos embed from grid."""
    if embed_dim % 2 != 0:
        raise ValueError("embed_dim must be even")

    omega = np.arange(embed_dim // 2, dtype=float)
    omega /= embed_dim / 2.0
    omega = 1.0 / 10000**omega

    pos = pos.reshape(-1)
    out = np.einsum("m,d->md", pos, omega)

    emb_sin = np.sin(out)
    emb_cos = np.cos(out)
    emb = np.concatenate([emb_sin, emb_cos], axis=1)
    return emb


class RAEDecoder(nn.Module):
    """RAE Decoder based on ViTMAE architecture.

    Takes latent features from the frozen encoder and reconstructs images.
    """

    def __init__(self, config: ViTMAEDecoderConfig, num_patches: int):
        super().__init__()
        self.config = config
        self.num_patches = num_patches

        # Decoder embedding: project from encoder hidden size to decoder hidden size
        self.decoder_embed = nn.Linear(config.hidden_size, config.decoder_hidden_size, bias=True)

        # Fixed sin-cos positional embedding
        self.decoder_pos_embed = nn.Parameter(
            torch.zeros(1, num_patches + 1, config.decoder_hidden_size), requires_grad=False
        )

        # Build decoder transformer layers
        decoder_config = ViTMAEDecoderConfig(
            hidden_size=config.decoder_hidden_size,
            num_hidden_layers=config.decoder_num_hidden_layers,
            num_attention_heads=config.decoder_num_attention_heads,
            intermediate_size=config.decoder_intermediate_size,
            hidden_act=config.hidden_act,
            hidden_dropout_prob=config.hidden_dropout_prob,
            attention_probs_dropout_prob=config.attention_probs_dropout_prob,
            initializer_range=config.initializer_range,
            layer_norm_eps=config.layer_norm_eps,
            image_size=config.image_size,
            patch_size=config.patch_size,
            num_channels=config.num_channels,
            qkv_bias=config.qkv_bias,
            decoder_hidden_size=config.decoder_hidden_size,
            decoder_num_attention_heads=config.decoder_num_attention_heads,
            decoder_num_hidden_layers=config.decoder_num_hidden_layers,
            decoder_intermediate_size=config.decoder_intermediate_size,
            mask_ratio=0.0,
            norm_pix_loss=False,
        )

        self.decoder_layers = nn.ModuleList(
            [ViTMAEDecoderLayer(decoder_config) for _ in range(config.decoder_num_hidden_layers)]
        )

        self.decoder_norm = nn.LayerNorm(config.decoder_hidden_size, eps=config.layer_norm_eps)
        self.decoder_pred = nn.Linear(
            config.decoder_hidden_size, config.patch_size**2 * config.num_channels, bias=True
        )

        # Trainable CLS token for decoder
        self.trainable_cls_token = nn.Parameter(torch.zeros(1, 1, config.decoder_hidden_size))

        self.initialize_weights(num_patches)

    def initialize_weights(self, num_patches: int) -> None:
        """Initialize positional embeddings and other weights."""
        # Initialize position embeddings with sin-cos
        decoder_pos_embed = get_2d_sincos_pos_embed(
            self.decoder_pos_embed.shape[-1], int(num_patches**0.5), add_cls_token=True
        )
        self.decoder_pos_embed.data.copy_(decoder_pos_embed.float().unsqueeze(0))

        # Initialize trainable CLS token
        nn.init.normal_(self.trainable_cls_token, std=self.config.initializer_range)

    def interpolate_latent(self, x: torch.Tensor) -> torch.Tensor:
        """Interpolate latent sequence to match expected number of patches."""
        b, l, c = x.shape
        if l == self.num_patches:
            return x

        # Interpolate latent to match num_patches
        h, w = int(l**0.5), int(l**0.5)
        x = x.reshape(b, h, w, c)
        x = x.permute(0, 3, 1, 2)
        target_size = (int(self.num_patches**0.5), int(self.num_patches**0.5))
        x = F.interpolate(x, size=target_size, mode="bilinear", align_corners=False)
        x = x.permute(0, 2, 3, 1).contiguous().view(b, self.num_patches, c)
        return x

    def unpatchify(self, patchified_pixel_values: torch.Tensor) -> torch.Tensor:
        """Convert patchified pixels back to image."""
        patch_size = self.config.patch_size
        num_channels = self.config.num_channels
        image_size = self.config.image_size

        num_patches_h = image_size // patch_size
        num_patches_w = image_size // patch_size

        batch_size = patchified_pixel_values.shape[0]
        patchified_pixel_values = patchified_pixel_values.reshape(
            batch_size,
            num_patches_h,
            num_patches_w,
            patch_size,
            patch_size,
            num_channels,
        )
        patchified_pixel_values = torch.einsum("nhwpqc->nchpwq", patchified_pixel_values)
        pixel_values = patchified_pixel_values.reshape(
            batch_size,
            num_channels,
            num_patches_h * patch_size,
            num_patches_w * patch_size,
        )
        return pixel_values

    def forward(
        self,
        hidden_states: torch.Tensor,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
    ) -> Union[RAEDecoderOutput, Tuple[torch.Tensor, ...]]:
        """Forward pass of the decoder.

        Args:
            hidden_states: Input latent from encoder, shape (B, N, C)
            output_attentions: Whether to output attention weights
            output_hidden_states: Whether to output all hidden states
            return_dict: Whether to return DecoderOutput or tuple

        Returns:
            RAEDecoderOutput with decoded sample or tuple of outputs
        """
        # Embed tokens (B, N, C) -> (B, N, decoder_hidden_size)
        x = self.decoder_embed(hidden_states)

        # Interpolate if needed
        x = self.interpolate_latent(x)

        # Add trainable CLS token
        cls_token = self.trainable_cls_token.expand(x.shape[0], -1, -1)
        x = torch.cat([cls_token, x], dim=1)

        # Add positional embeddings
        hidden_states = x + self.decoder_pos_embed

        # Apply transformer layers
        all_hidden_states = () if output_hidden_states else None
        all_self_attentions = () if output_attentions else None

        for layer_module in self.decoder_layers:
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_outputs = layer_module(hidden_states, output_attentions=output_attentions)
            hidden_states = layer_outputs[0]

            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_outputs[1],)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        # Final norm
        hidden_states = self.decoder_norm(hidden_states)

        # Predict pixel values
        logits = self.decoder_pred(hidden_states)

        # Remove CLS token
        logits = logits[:, 1:, :]

        # Unpatchify to image
        sample = self.unpatchify(logits)

        if not return_dict:
            return (sample,)

        return RAEDecoderOutput(sample=sample)


class AutoencoderRAE(
    ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapterMixin
):
    r"""
    Representation Autoencoder (RAE) model for encoding images to latents and decoding latents to images.

    This model uses a frozen pretrained encoder (DINOv2, SigLIP2, or MAE) with a trainable ViT decoder
    to reconstruct images from learned representations.

    This model inherits from [`ModelMixin`]. Check the superclass documentation for its generic methods
    implemented for all models (such as downloading or saving).

    Args:
        encoder_type (`str`, *optional*, defaults to `"dinov2"`):
            Type of frozen encoder to use. One of `"dinov2"`, `"siglip2"`, or `"mae"`.
        encoder_name_or_path (`str`, *optional*, defaults to `"facebook/dinov2-with-registers-base"`):
            Path to pretrained encoder model or model identifier from huggingface.co/models.
        decoder_config (`ViTMAEDecoderConfig`, *optional*):
            Configuration for the decoder. If None, a default config will be used.
        num_patches (`int`, *optional*, defaults to `196`):
            Number of patches in the latent space (14x14 = 196 for 224x224 image with patch size 16).
        patch_size (`int`, *optional*, defaults to `16`):
            Patch size for both encoder and decoder.
        encoder_input_size (`int`, *optional*, defaults to `224`):
            Input size expected by the encoder.
        image_size (`int`, *optional*, defaults to `256`):
            Output image size.
        num_channels (`int`, *optional*, defaults to `3`):
            Number of input/output channels.
        latent_mean (`torch.Tensor`, *optional*):
            Optional mean for latent normalization.
        latent_var (`torch.Tensor`, *optional*):
            Optional variance for latent normalization.
        noise_tau (`float`, *optional*, defaults to `0.0`):
            Noise level for training (adds noise to latents during training).
        reshape_to_2d (`bool`, *optional*, defaults to `True`):
            Whether to reshape latents to 2D (B, C, H, W) format.
        use_encoder_loss (`bool`, *optional*, defaults to `False`):
            Whether to use encoder hidden states in the loss (for advanced training).
    """

    _supports_gradient_checkpointing = True
    _no_split_modules = ["ViTMAEDecoderLayer"]

    @register_to_config
    def __init__(
        self,
        encoder_type: str = "dinov2",
        encoder_name_or_path: str = "facebook/dinov2-with-registers-base",
        decoder_config: Optional[ViTMAEDecoderConfig] = None,
        num_patches: int = 196,
        patch_size: int = 16,
        encoder_input_size: int = 224,
        image_size: int = 256,
        num_channels: int = 3,
        latent_mean: Optional[torch.Tensor] = None,
        latent_var: Optional[torch.Tensor] = None,
        noise_tau: float = 0.0,
        reshape_to_2d: bool = True,
        use_encoder_loss: bool = False,
    ):
        super().__init__()

        self.encoder_type = encoder_type
        self.encoder_name_or_path = encoder_name_or_path
        self.patch_size = patch_size
        self.encoder_input_size = encoder_input_size
        self.image_size = image_size
        self.num_channels = num_channels
        self.num_patches = num_patches
        self.noise_tau = noise_tau
        self.reshape_to_2d = reshape_to_2d
        self.use_encoder_loss = use_encoder_loss

        # Default decoder config if not provided
        if decoder_config is None:
            decoder_config = ViTMAEDecoderConfig(
                hidden_size=768,  # Match DINOv2-B
                num_hidden_layers=12,
                num_attention_heads=12,
                intermediate_size=3072,
                decoder_hidden_size=768,
                decoder_num_attention_heads=12,
                decoder_num_hidden_layers=12,
                decoder_intermediate_size=3072,
                patch_size=patch_size,
                image_size=image_size,
                num_channels=num_channels,
            )
        elif isinstance(decoder_config, dict):
            # Convert dict to ViTMAEDecoderConfig during loading
            decoder_config = ViTMAEDecoderConfig(**decoder_config)

        self.decoder_config = decoder_config
        self.decoder = RAEDecoder(decoder_config, num_patches)

        # Encoder will be initialized lazily in encode()
        self._encoder = None
        self._encoder_processor = None

        # Latent normalization - use internal names to avoid conflict with config
        self.register_buffer("_latent_mean", latent_mean, persistent=False)
        self.register_buffer("_latent_var", latent_var, persistent=False)
        self.do_normalization = latent_mean is not None
        self.eps = 1e-5

        # Encoder normalization stats (set during encode based on encoder type)
        self.register_buffer("encoder_mean", None, persistent=False)
        self.register_buffer("encoder_std", None, persistent=False)

    def __getattr__(self, name: str):
        """Handle decoder_config serialization by converting to dict when accessed via config."""
        # Check if this is a decoder_config attribute being accessed for serialization
        if name == "decoder_config" and name not in self.__dict__:
            # Return the config as a dict for serialization purposes
            return self._decoder_config_dict
        return super().__getattr__(name)

    @property
    def _decoder_config_dict(self) -> dict:
        """Return decoder_config as dict for serialization."""
        if hasattr(self.decoder_config, 'to_dict'):
            return self.decoder_config.to_dict()
        return vars(self.decoder_config)

    @property
    def encoder(self):
        """Lazy load the encoder."""
        if self._encoder is None:
            self._init_encoder()
        return self._encoder

    @property
    def encoder_processor(self):
        """Lazy load the encoder processor."""
        if self._encoder_processor is None:
            self._init_encoder()
        return self._encoder_processor

    def _init_encoder(self):
        """Initialize the encoder and processor based on encoder_type."""
        from transformers import AutoImageProcessor

        if self.encoder_type == "dinov2":
            from transformers import Dinov2WithRegistersModel

            self._encoder = Dinov2WithRegistersModel.from_pretrained(self.encoder_name_or_path)
            self._encoder.requires_grad_(False)
            # Remove layernorm affine parameters
            self._encoder.layernorm.elementwise_affine = False
            self._encoder.layernorm.weight = None
            self._encoder.layernorm.bias = None

            proc = AutoImageProcessor.from_pretrained(self.encoder_name_or_path)
            self.encoder_mean = torch.tensor(proc.image_mean).view(1, 3, 1, 1)
            self.encoder_std = torch.tensor(proc.image_std).view(1, 3, 1, 1)

            self.patch_size = self._encoder.config.patch_size
            self.hidden_size = self._encoder.config.hidden_size
            # Sync decoder config hidden_size with encoder after loading
            self.decoder_config.hidden_size = self.hidden_size
            # Update decoder_embed input dimension to match encoder
            self.decoder.decoder_embed = nn.Linear(
                self.hidden_size, self.decoder_config.decoder_hidden_size, bias=True
            )

        elif self.encoder_type == "siglip2":
            from transformers import SiglipModel

            self._encoder = SiglipModel.from_pretrained(self.encoder_name_or_path)
            self._encoder.requires_grad_(False)
            # Remove post_layernorm affine parameters
            self._encoder.vision_model.post_layernorm.elementwise_affine = False
            self._encoder.vision_model.post_layernorm.weight = None
            self._encoder.vision_model.post_layernorm.bias = None

            proc = AutoImageProcessor.from_pretrained(self.encoder_name_or_path)
            self.encoder_mean = torch.tensor(proc.image_mean).view(1, 3, 1, 1)
            self.encoder_std = torch.tensor(proc.image_std).view(1, 3, 1, 1)

            self.patch_size = self._encoder.config.patch_size
            self.hidden_size = self._encoder.config.hidden_size
            # Sync decoder config hidden_size with encoder after loading
            self.decoder_config.hidden_size = self.hidden_size
            # Update decoder_embed input dimension to match encoder
            self.decoder.decoder_embed = nn.Linear(
                self.hidden_size, self.decoder_config.decoder_hidden_size, bias=True
            )

        elif self.encoder_type == "mae":
            from transformers import ViTMAEForPreTraining

            vit_mae = ViTMAEForPreTraining.from_pretrained(self.encoder_name_or_path)
            self._encoder = vit_mae.vit
            self._encoder.requires_grad_(False)
            # Remove layernorm affine parameters
            self._encoder.layernorm.elementwise_affine = False
            self._encoder.layernorm.weight = None
            self._encoder.layernorm.bias = None

            self._encoder.config.mask_ratio = 0.0  # No masking

            proc = AutoImageProcessor.from_pretrained(self.encoder_name_or_path)
            self.encoder_mean = torch.tensor(proc.image_mean).view(1, 3, 1, 1)
            self.encoder_std = torch.tensor(proc.image_std).view(1, 3, 1, 1)

            self.patch_size = self._encoder.config.patch_size
            self.hidden_size = self._encoder.config.hidden_size
            # Sync decoder config hidden_size with encoder after loading
            self.decoder_config.hidden_size = self.hidden_size
            # Update decoder_embed input dimension to match encoder
            self.decoder.decoder_embed = nn.Linear(
                self.hidden_size, self.decoder_config.decoder_hidden_size, bias=True
            )

        else:
            raise ValueError(f"Unknown encoder_type: {self.encoder_type}. Supported: dinov2, siglip2, mae")

    def _encode_dinov2(self, x: torch.Tensor) -> torch.Tensor:
        """Encode using DINOv2."""
        outputs = self._encoder(x, output_hidden_states=True)
        # Remove CLS token and register tokens (5 tokens total: 1 CLS + 4 registers)
        image_features = outputs.last_hidden_state[:, 5:]
        return image_features

    def _encode_siglip2(self, x: torch.Tensor) -> torch.Tensor:
        """Encode using SigLIP2."""
        outputs = self._encoder.vision_model(x, output_hidden_states=True, interpolate_pos_encoding=True)
        return outputs.last_hidden_state

    def _encode_mae(self, x: torch.Tensor) -> torch.Tensor:
        """Encode using MAE."""
        h, w = x.shape[2], x.shape[3]
        patch_num = int(h * w // self.patch_size**2)
        noise = torch.arange(patch_num).unsqueeze(0).expand(x.shape[0], -1).to(x.device).to(x.dtype)
        outputs = self._encoder(x, noise, interpolate_pos_encoding=True)
        # Remove CLS token
        image_features = outputs.last_hidden_state[:, 1:]
        return image_features

    def noising(self, x: torch.Tensor) -> torch.Tensor:
        """Add noise to latents (for training)."""
        noise_sigma = self.noise_tau * torch.rand((x.size(0),) + (1,) * (len(x.shape) - 1), device=x.device)
        noise = noise_sigma * torch.randn_like(x)
        return x + noise

    @apply_forward_hook
    def encode(
        self, x: torch.Tensor, return_dict: bool = True
    ) -> Union[AutoencoderRAEOuput, Tuple[torch.Tensor, ...]]:
        """
        Encode a batch of images into latents.

        Args:
            x (`torch.Tensor`): Input batch of images. Shape: (B, C, H, W)
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether to return a [`~AutoencoderRAEOuput`] instead of a plain tuple.

        Returns:
            [`~AutoencoderRAEOuput`] or `tuple`:
                If return_dict is True, a [`~AutoencoderRAEOuput`] is returned, otherwise a plain `tuple` is returned.
                The latent is a tensor of shape (B, hidden_size, latent_h, latent_w).
        """
        batch_size, num_channels, height, width = x.shape

        # Resize input if needed
        if height != self.encoder_input_size or width != self.encoder_input_size:
            x = F.interpolate(
                x,
                size=(self.encoder_input_size, self.encoder_input_size),
                mode="bicubic",
                align_corners=False,
            )

        # Initialize encoder if not already done (handles lazy loading after from_pretrained)
        if self._encoder is None:
            self._init_encoder()

        # Normalize input
        x = (x - self.encoder_mean.to(x.device)) / self.encoder_std.to(x.device)

        # Encode based on encoder type
        if self.encoder_type == "dinov2":
            z = self._encode_dinov2(x)
        elif self.encoder_type == "siglip2":
            z = self._encode_siglip2(x)
        elif self.encoder_type == "mae":
            z = self._encode_mae(x)
        else:
            raise ValueError(f"Unknown encoder_type: {self.encoder_type}")

        # Add noise during training
        if self.training and self.noise_tau > 0:
            z = self.noising(z)

        # Reshape to 2D if needed
        if self.reshape_to_2d:
            b, n, c = z.shape
            h = w = int(sqrt(n))
            z = z.transpose(1, 2).view(b, c, h, w)

        # Apply latent normalization if enabled
        if self.do_normalization:
            latent_mean = self._latent_mean.to(z.device) if self._latent_mean is not None else 0
            latent_var = self._latent_var.to(z.device) if self._latent_var is not None else 1
            z = (z - latent_mean) / torch.sqrt(latent_var + self.eps)

        if not return_dict:
            return (z,)

        return AutoencoderRAEOuput(latent=z)

    def _decode(self, z: torch.Tensor) -> RAEDecoderOutput:
        """Decode latents to images (internal method)."""
        # Apply inverse latent normalization if enabled
        if self.do_normalization:
            latent_mean = self._latent_mean.to(z.device) if self._latent_mean is not None else 0
            latent_var = self._latent_var.to(z.device) if self._latent_var is not None else 1
            z = z * torch.sqrt(latent_var + self.eps) + latent_mean

        # Reshape from 2D to sequence if needed
        if self.reshape_to_2d:
            b, c, h, w = z.shape
            n = h * w
            # Properly reshape from (B, C, H, W) to (B, N, C)
            z = z.reshape(b, c, n)  # (B, C, N)
            z = z.permute(0, 2, 1)  # (B, N, C)

        # Decode
        decoder_output = self.decoder(z, return_dict=True)
        sample = decoder_output.sample

        # Denormalize output
        sample = sample * self.encoder_std.to(sample.device) + self.encoder_mean.to(sample.device)

        return RAEDecoderOutput(sample=sample)

    @apply_forward_hook
    def decode(
        self, z: torch.Tensor, return_dict: bool = True
    ) -> Union[RAEDecoderOutput, torch.Tensor]:
        """
        Decode a batch of latents into images.

        Args:
            z (`torch.Tensor`): Input batch of latents. Shape: (B, C, H, W)
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether to return a [`~RAEDecoderOutput`] instead of a plain tensor.

        Returns:
            [`~RAEDecoderOutput`] or `torch.Tensor`:
                If return_dict is True, a [`~RAEDecoderOutput`] is returned, otherwise a plain tensor is returned.
        """
        decoder_output = self._decode(z)

        if not return_dict:
            return decoder_output.sample

        return decoder_output

    def forward(
        self,
        sample: torch.Tensor,
        return_dict: bool = True,
    ) -> Union[RAEDecoderOutput, torch.Tensor]:
        """
        Forward pass: encode then decode.

        Args:
            sample (`torch.Tensor`): Input image tensor. Shape: (B, C, H, W)
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether to return a [`~RAEDecoderOutput`] instead of a plain tensor.

        Returns:
            [`~RAEDecoderOutput`] or `torch.Tensor`:
                Reconstructed image tensor. Shape: (B, C, H, W)
        """
        # Encode
        encode_output = self.encode(sample, return_dict=True)
        z = encode_output.latent

        # Decode
        decode_output = self.decode(z, return_dict=True)

        if not return_dict:
            return decode_output.sample

        return RAEDecoderOutput(sample=decode_output.sample)

RAE support #13000

Description

Metadata

Metadata

Assignees

Labels

Type

Fields

Projects

Milestone

Relationships

Development

Issue actions