feat: implement rae autoencoder. by Ando233 · Pull Request #13046 · huggingface/diffusers

Ando233 · 2026-01-28T12:31:05Z

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 · 2026-01-30T11:32:07Z

@bytetriper if you could take a look?

kashif · 2026-01-30T14:26:23Z

nice works @Ando233 checking

kashif · 2026-01-30T14:47:34Z

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 · 2026-01-31T01:51:23Z

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 · 2026-01-31T03:29:30Z

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

kashif · 2026-02-15T23:39:12Z

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

kashif · 2026-02-15T23:40:25Z

@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?

HuggingFaceDocBuilderDev · 2026-02-27T10:31:02Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

sayakpaul

Left some comments. Let me know if this makes sense. @bytetriper it would be great if you could also test the diffusers counterparts of RAE and let us know your thoughts.

docs/source/en/api/models/autoencoder_rae.md

examples/research_projects/autoencoder_rae/README.md

src/diffusers/models/autoencoders/autoencoder_rae.py

tests/models/autoencoders/test_models_autoencoder_rae.py

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

sayakpaul

Left a major comment regarding the presence of the encoder-specific classes now. LMK your thoughts.

tests/models/autoencoders/test_models_autoencoder_rae.py

sayakpaul · 2026-03-02T05:18:52Z

src/diffusers/models/autoencoders/autoencoder_rae.py

+        self.model.layernorm.weight = None
+        self.model.layernorm.bias = None


We're already stripping the layernorms in the conversion. Seems like it's not needed anymore?

sayakpaul · 2026-03-02T05:19:51Z

src/diffusers/models/autoencoders/autoencoder_rae.py

+        self.model.vision_model.post_layernorm.weight = None
+        self.model.vision_model.post_layernorm.bias = None


Same for these?

src/diffusers/models/autoencoders/autoencoder_rae.py

kashif · 2026-03-03T09:04:10Z

@bytetriper could you kindly merge the change in the PR on the hub for the weights?

sayakpaul

Looks really good barring a few nits.

sayakpaul · 2026-03-03T11:30:04Z

src/diffusers/models/autoencoders/autoencoder_rae.py

+        # RAE-main: base_patches = (encoder_input_size // encoder_patch_size) ** 2
+        if self.encoder_input_size % encoder_patch_size != 0:
+            raise ValueError(
+                f"encoder_input_size={self.encoder_input_size} must be divisible by encoder_patch_size={encoder_patch_size}."
+            )


Let's error as early as possible?

sayakpaul · 2026-03-03T11:31:50Z

src/diffusers/models/autoencoders/autoencoder_rae.py

+        model.layernorm.weight = None
+        model.layernorm.bias = None


Can we add a small comment explaining what this is doing?

sayakpaul · 2026-03-03T11:32:38Z

src/diffusers/models/autoencoders/autoencoder_rae.py

+        # Decoder patch size is independent from encoder patch size.
+        decoder_patch_size = int(patch_size)
+        if decoder_patch_size <= 0:
+            raise ValueError("patch_size must be a positive integer (this is decoder_patch_size).")


Same. Let's error as early as possible.

sayakpaul · 2026-03-03T11:34:44Z

tests/models/autoencoders/test_models_autoencoder_rae.py

+
+@slow
+@pytest.mark.skip(reason="Not enough model usage to justify slow tests yet.")
+class AutoencoderRAEEncoderIntegrationTests(unittest.TestCase):


We shouldn't inherit from unittest now that we're using pure pytest.

sayakpaul · 2026-03-03T11:35:44Z

tests/models/autoencoders/test_models_autoencoder_rae.py

+        gc.collect()
+        backend_empty_cache(torch_device)
+
+    def test_autoencoder_rae_from_pretrained_dinov2(self):


Should we check against an image and assert a value slice?

sayakpaul · 2026-03-03T11:36:35Z

@dg845 could you give this a review as well?

bytetriper · 2026-03-03T18:52:09Z

@bytetriper could you kindly merge the change in the PR on the hub for the weights?

@kashif Merged!

dg845 · 2026-03-04T07:08:53Z

src/diffusers/models/autoencoders/autoencoder_rae.py

+def _build_encoder(encoder_type: str, hidden_size: int, patch_size: int, num_hidden_layers: int) -> nn.Module:
+    """Build a frozen encoder from config (no pretrained download)."""
+    num_attention_heads = hidden_size // 64  # all supported encoders use head_dim=64


Suggested change

def _build_encoder(encoder_type: str, hidden_size: int, patch_size: int, num_hidden_layers: int) -> nn.Module:

"""Build a frozen encoder from config (no pretrained download)."""

num_attention_heads = hidden_size // 64 # all supported encoders use head_dim=64

def _build_encoder(

encoder_type: str, hidden_size: int, patch_size: int, num_hidden_layers: int, head_dim: int = 64

) -> nn.Module:

"""Build a frozen encoder from config (no pretrained download)."""

num_attention_heads = hidden_size // head_dim # all supported encoders use head_dim=64

nit: I think having a head_dim argument here would be useful to make the code more future proof (for example, if a CoolNewRepresentatonEncoder comes out in the future with head_dim != 64, it would be easier to support).

dg845 · 2026-03-04T07:12:11Z

src/diffusers/models/autoencoders/autoencoder_rae.py

+        self.decoder_pos_embed = nn.Parameter(
+            torch.zeros(1, num_patches + 1, decoder_hidden_size), requires_grad=False
+        )


Is setting requires_grad=False here intended? If decoder_pos_embed is meant to be a fixed 2D sinusoidal positional embedding (as I believe _initialize_weights is doing below), perhaps we could register it as a buffer instead?

dg845 · 2026-03-04T07:15:36Z

src/diffusers/models/autoencoders/autoencoder_rae.py

+        self.gradient_checkpointing = False
+
+        self._initialize_weights(num_patches)
+        self.set_trainable_cls_token()


Should set_trainable_cls_token receive an argument from __init__ here? If not, maybe it would be more clear to inline its logic, e.g.

self.trainable_cls_token = nn.Parameter(torch.zeros(1, 1, self.decoder_hidden_size))

dg845 · 2026-03-04T07:21:45Z

src/diffusers/models/autoencoders/autoencoder_rae.py

+        x = x.permute(0, 2, 3, 1).contiguous().view(b, self.num_patches, c)
+        return x
+
+    def unpatchify(self, patchified_pixel_values: torch.Tensor, original_image_size: Optional[Tuple[int, int]] = None):


Suggested change

def unpatchify(self, patchified_pixel_values: torch.Tensor, original_image_size: Optional[Tuple[int, int]] = None):

def unpatchify(self, patchified_pixel_values: torch.Tensor, original_image_size: tuple[int, int] | None = None):

We have moved to Python 3.9+ style implicit type annotations. Can you update the type annotations here and elsewhere where the old style is used? See #12524 for more details.

dg845 · 2026-03-04T07:23:18Z

src/diffusers/models/autoencoders/autoencoder_rae.py

+        self._initialize_weights(num_patches)
+        self.set_trainable_cls_token()
+
+    def set_trainable_cls_token(self, tensor: Optional[torch.Tensor] = None):


Suggested change

def set_trainable_cls_token(self, tensor: Optional[torch.Tensor] = None):

def set_trainable_cls_token(self, tensor: torch.Tensor | None = None):

See #13046.

dg845 · 2026-03-04T07:28:46Z

src/diffusers/models/autoencoders/autoencoder_rae.py

+        # Slicing support (batch dimension) similar to other diffusers autoencoders
+        self.use_slicing = False
+
+    def _noising(self, x: torch.Tensor) -> torch.Tensor:


Could the _noising method take a generator argument for reproducibility?

dg845 · 2026-03-04T07:31:52Z

src/diffusers/models/autoencoders/autoencoder_rae.py

+        noise_sigma = self.noise_tau * torch.rand((x.size(0),) + (1,) * (x.ndim - 1), device=x.device, dtype=x.dtype)
+        return x + noise_sigma * torch.randn_like(x)
+
+    def _maybe_resize_and_normalize(self, x: torch.Tensor) -> torch.Tensor:


Suggested change

def _maybe_resize_and_normalize(self, x: torch.Tensor) -> torch.Tensor:

def _resize_and_normalize(self, x: torch.Tensor) -> torch.Tensor:

nit: I think it would be more clear to drop the _maybe prefixes from the _maybe_* methods as it looks like all of these methods unconditionally perform the action in the name.

wangyuqi and others added 3 commits January 25, 2026 02:54

feat: implement three RAE encoders(dinov2, siglip2, mae)

382aad0

feat: finish first version of autoencoder_rae

f82cecc

Merge branch 'main' into rae

a3926d7

sayakpaul requested a review from kashif January 30, 2026 11:31

kashif added 9 commits February 15, 2026 23:05

Merge branch 'main' into rae

3ecf89d

fix formatting

0850c8c

make fix-copies

24acab0

initial doc

25bc9e3

fix latent_mean / latent_var init types to accept config-friendly inputs

f06ea7a

use mean and std convention

d7cb124

cleanup

0d59b22

add rae to diffusers script

202b14f

use imports

7cbbf27

kashif added 10 commits February 15, 2026 23:50

use attention

e6d4499

remove unneeded class

6a9bde6

example traiing script

9522e68

input and ground truth sizes have to be the same

906d79a

fix argument

d3cbd5a

move loss to training script

96520c4

cleanup

fc52959

simplify mixins

a4fc9f6

fix training script

d06b501

Merge branch 'main' into rae

d8b2983

sayakpaul reviewed Feb 27, 2026

View reviewed changes

kashif and others added 9 commits February 27, 2026 11:42

Update examples/research_projects/autoencoder_rae/README.md

6a78767

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

Update src/diffusers/models/autoencoders/autoencoder_rae.py

1b4a43f

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

use image url

c717498

Encoder is frozen

03e757c

fix slow test

8861a80

remove config

9ada576

use ModelTesterMixin and AutoencoderTesterMixin

34787e5

make quality

cd86873

Merge branch 'main' into rae

bff48d3

sayakpaul requested a review from stevhliu February 28, 2026 16:49

kashif added 2 commits February 28, 2026 19:35

strip final layernorm when converting

f15873a

_strip_final_layernorm_affine for training script

e834e49

sayakpaul reviewed Mar 2, 2026

View reviewed changes

kashif added 2 commits March 2, 2026 08:11

fix test

5817416

add dispatch forward and update conversion script

3a66895

update training script

bc71889

sayakpaul reviewed Mar 3, 2026

View reviewed changes

Merge branch 'main' into rae

9c0f96b

dg845 reviewed Mar 4, 2026

View reviewed changes

		self.model.layernorm.weight = None
		self.model.layernorm.bias = None

		self.model.vision_model.post_layernorm.weight = None
		self.model.vision_model.post_layernorm.bias = None

	def unpatchify(self, patchified_pixel_values: torch.Tensor, original_image_size: Optional[Tuple[int, int]] = None):
	def unpatchify(self, patchified_pixel_values: torch.Tensor, original_image_size: tuple[int, int] \| None = None):

	def set_trainable_cls_token(self, tensor: Optional[torch.Tensor] = None):
	def set_trainable_cls_token(self, tensor: torch.Tensor \| None = None):

	def _maybe_resize_and_normalize(self, x: torch.Tensor) -> torch.Tensor:
	def _resize_and_normalize(self, x: torch.Tensor) -> torch.Tensor:

Conversation

Ando233 commented Jan 28, 2026

What does this PR do?

Before submitting

Usage

Who can review?

Uh oh!

sayakpaul commented Jan 30, 2026

Uh oh!

kashif commented Jan 30, 2026

Uh oh!

kashif commented Jan 30, 2026

Uh oh!

bytetriper commented Jan 31, 2026

Uh oh!

sayakpaul commented Jan 31, 2026

Uh oh!

kashif commented Feb 15, 2026

Uh oh!

kashif commented Feb 15, 2026

Uh oh!

HuggingFaceDocBuilderDev commented Feb 27, 2026

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

kashif commented Mar 3, 2026

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul commented Mar 3, 2026

Uh oh!

bytetriper commented Mar 3, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dg845 Mar 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dg845 Mar 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dg845 Mar 4, 2026 •

edited

Loading

dg845 Mar 4, 2026 •

edited

Loading