feat: implement rae autoencoder.

[Ando233]

What does this PR do?

This PR adds a new representation autoencoder implementation, AutoencoderRAE, to diffusers.
Implements diffusers.models.autoencoders.autoencoder_rae.AutoencoderRAE with a frozen pretrained vision encoder (DINOv2 / SigLIP2 / ViT-MAE) and a ViT-MAE style decoder.
The decoder implementation is aligned with the RAE-main GeneralDecoder parameter structure, enabling loading of existing trained decoder checkpoints (e.g. model.pt) without key mismatches when encoder/decoder settings are consistent.
Adds unit/integration tests under diffusers/tests/models/autoencoders/test_models_autoencoder_rae.py.
Registers exports so users can import directly via from diffusers import AutoencoderRAE.

Fixes #13000

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case: RAE support #13000
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Usage

ae = AutoencoderRAE(
    encoder_cls="dinov2",
    encoder_name_or_path=encoder_path,
    image_size=image_size,
    encoder_input_size=image_size,
    patch_size=patch_size,
    num_patches=num_patches,
    decoder_hidden_size=1152,
    decoder_num_hidden_layers=28,
    decoder_num_attention_heads=16,
    decoder_intermediate_size=4096,
).to(device)
ae.eval()

state = torch.load(args.decoder_ckpt, map_location="cpu")
ae.decoder.load_state_dict(state, strict=False)

with torch.no_grad():
    recon = ae(x).sample

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[sayakpaul]

@bytetriper if you could take a look?

[kashif]

nice works @Ando233 checking

[kashif]

off the bat,

• let's have a nice convention for the output datatype classes, have a look at the other autoencoder for the convention in difusers
• some of the tests might need to be marked as slow and some paths are hard-coded

lets sort out these things and then re-look

[bytetriper]

Agree with @kashif . Also if possible we can bake all the params into config so we can enable .from_pretrained(), which is more elegant and aligns with diffusers usage. I can help convert our released ckpt to hgf format afterwards

[sayakpaul]

@Ando233 we're happy to provide assistance if needed.

[kashif]

@Ando233 the one remaining thing is the use of the use_encoder_loss and perhaps an example real-world training script

[kashif]

@bytetriper could you kindly try to run the conversion scripts and upload the diffusers style weights to your huggingface hub for the checkpoints you have?

[kashif]

@dg845 resolved issues thanks

[sayakpaul]

@Ando233 / @kashif could you merge Ando233#1?

[sayakpaul]

What happens if we directly initialize this to what we're doing in the initialize_weights() function? Could we get rid of the explicit device placement in return x_rec.to(device=z.device) then?

[sayakpaul]

I guess a simpler solution would be to directly assign the pos_embed value (we are initializing through initalize_weights()) and just persist it in the state dict. That way, we can skip explicit device placements like return x_rec.to(device=z.device)?

This would require opening PRs to the RAE repos on the Hub, though.

@dg845 WDYT?

[dg845]

I think something like

class RAEDecoder(nn.Module):
    ...
    def __init__(...):
        ...
        grid_size = int(num_patches**0.5)
        pos_embed = get_2d_sincos_pos_embed(
            decoder_hidden_size,
            grid_size,
            cls_token=True,
            extra_tokens=1,
            output_type="pt",
        )
        self.register_buffer("decoder_pos_embed", pos_embed.unsqueeze(0), persistent=False)
        ...

is generally how we would implement this, see e.g. here:

diffusers/src/diffusers/models/transformers/transformer_wan.py

Lines 392 to 393 in ab6040a

	self.register_buffer("freqs_cos", torch.cat(freqs_cos, dim=1), persistent=False)
	self.register_buffer("freqs_sin", torch.cat(freqs_sin, dim=1), persistent=False)

I think this should handle device placement automatically and also avoid the need to change the Hub repo (although maybe the conversion script might need to be changed?).

[dg845]

For test_model_parallelism specifically, the error I encountered earlier in #13046 (comment) may be the result of an accelerate bug where the decoder_pos_embed buffer doesn't end up on the device_map, so AutoencoderRAE.from_pretrained(..., device_map="auto") doesn't know where to put it and gives that error. I've opened an issue for this at huggingface/accelerate#3956.

[sayakpaul]

@kashif WDYT? Could we make this change?

[kashif]

yes i can do it

[sayakpaul]

Thanks! Then let's prefer thi solution.

[sayakpaul]

Thanks a lot @kashif for shipping RAEs!