diffusers/pr_13231/en/api/models/rae_dit_transformer2d.md

RAEDiT2DModel

The RAEDiT2DModel is the Stage-2 latent diffusion transformer introduced in Diffusion Transformers with Representation Autoencoders.

Unlike DiT models that operate on VAE latents, this transformer denoises the latent space learned by AutoencoderRAE. It is designed to be used with FlowMatchEulerDiscreteScheduler and decoded back to RGB with AutoencoderRAE.

Loading a pretrained transformer

from diffusers import RAEDiT2DModel

transformer = RAEDiT2DModel.from_pretrained("path/to/converted-stage2-transformer")

RAEDiT2DModel[[diffusers.RAEDiT2DModel]]

diffusers.RAEDiT2DModel[[diffusers.RAEDiT2DModel]]

Source

Stage-2 latent diffusion transformer used by the RAE paper.

The architecture mirrors the upstream two-stream DiTwDDTHead design: an encoder path first builds conditioning tokens from the latent input, then a decoder path denoises the latent tokens conditioned on those encoded tokens.