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?

Ando233 · 2026-02-17T15:06:00Z

Thank you for efforts @kashif , let me try to implement the remaining use_encoder_loss and real-world training script

kashif · 2026-02-17T16:15:59Z

@Ando233 I added that already, so next we can wait for @bytetriper for a review and see if the weight conversion works on his end

bytetriper · 2026-02-22T21:27:33Z

Thanks for the implementation! I just checked and weight conversion works on my end. Converted models are under https://huggingface.co/collections/nyu-visionx/rae. @kashif @Ando233 Can you check whether the converted models work on your end?

sayakpaul · 2026-02-23T02:50:30Z

@bytetriper thanks! What would be the quickest way to validate if the implementation is correct? We can do a quick value assertion test between the original model and the converted model on the same inputs. Would you be able to do it?

sayakpaul

Left a bunch of comments. The major thing is we need to be a bit more explicit in terms of how we're defining the configs, loading encoder state dicts, etc.

I think we could aim for the following entrypoint for instantiating the AutoencoderRAE class:

AutoencoderRAE(..., encoder_type="dinov2")

Inside the implementation of AutoencoderRAE __init__(), specifically, we can have a simple if/else block to dispatch the encoder based on encoder_type.:

if encoder_type == "dinov2":
    encoder = Dinov2Encoder()
elif encoder_type == "siglip2":
    encoder = Siglip2Encoder()
...

And then, when a user does AutoencoderRAE.from_pretrained(...), the state dict should have both the encoder and decoder state dict, following how it's done in the other Autoencoder implementations of diffusers.

I will also let @dg845 take a look and provide feedback.

sayakpaul · 2026-02-23T02:52:33Z

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

+`AutoencoderRAE` is a representation autoencoder that combines a frozen vision encoder (DINOv2, SigLIP2, or MAE) with a ViT-MAE-style decoder.
+
+Paper: [Diffusion Transformers with Representation Autoencoders](https://huggingface.co/papers/2510.11690).
+
+The model follows the standard diffusers autoencoder API:
+- `encode(...)` returns an `EncoderOutput` with a `latent` tensor.
+- `decode(...)` returns a `DecoderOutput` with a `sample` tensor.


Cc: @stevhliu. Could you leave suggestions on the docs?

sayakpaul · 2026-02-23T02:53:39Z

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

+
+model = AutoencoderRAE(
+    encoder_cls="dinov2",
+    encoder_name_or_path="facebook/dinov2-with-registers-base",


Org should nyu-visionx.

sayakpaul · 2026-02-23T02:53:57Z

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

+- `encode(...)` returns an `EncoderOutput` with a `latent` tensor.
+- `decode(...)` returns a `DecoderOutput` with a `sample` tensor.
+
+## Usage


@kashif does this need updating?

sayakpaul · 2026-02-23T02:54:50Z

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

+
+For latent normalization, use `latents_mean` and `latents_std` (matching other diffusers autoencoders).
+
+See `examples/research_projects/autoencoder_rae/train_autoencoder_rae.py` for a stage-1 style training script


What does stage-2 have? Generation?

sayakpaul · 2026-02-23T02:55:10Z

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

+
+`encoder_cls` supports `"dinov2"`, `"siglip2"`, and `"mae"`.
+
+For latent normalization, use `latents_mean` and `latents_std` (matching other diffusers autoencoders).


We should provide an example for this.

sayakpaul · 2026-02-23T03:09:29Z

src/diffusers/models/autoencoders/autoencoder_rae.py

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


These params are not used in the forward pass anyway. So, maybe it's not needed?

This seems to have been not addressed.

sayakpaul · 2026-02-23T03:11:31Z

src/diffusers/models/autoencoders/autoencoder_rae.py

+        from transformers import AutoImageProcessor
+
+        proc = AutoImageProcessor.from_pretrained(encoder_name_or_path)
+        encoder_mean = torch.tensor(proc.image_mean, dtype=torch.float32).view(1, 3, 1, 1)
+        encoder_std = torch.tensor(proc.image_std, dtype=torch.float32).view(1, 3, 1, 1)


This should be explicitly in the conversion script. This is an antipattern for the library.

We could do something like:
https://github.com/huggingface/diffusers/blob/a80b19218b4bd4faf2d6d8c428dcf1ae6f11e43d/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py#L1112C9-L1116C1

Then in the conversion script, make these a part of the converted state dict before loading that into the diffusers implementation. LMK if it's unclear.

sayakpaul · 2026-02-23T03:15:42Z

src/diffusers/models/autoencoders/autoencoder_rae.py

+        # Optional latent normalization (RAE-main uses mean/var)
+        latents_mean_tensor = _as_optional_tensor(latents_mean)
+        self.do_latent_normalization = latents_mean is not None or latents_std is not None
+        if latents_mean_tensor is not None:
+            self.register_buffer("_latents_mean", latents_mean_tensor, persistent=True)
+        else:
+            self._latents_mean = None
+        if latents_std_tensor is not None:
+            self.register_buffer("_latents_std", latents_std_tensor, persistent=True)
+        else:
+            self._latents_std = None


Seems like this can be removed?

sayakpaul · 2026-02-23T03:15:56Z

src/diffusers/models/autoencoders/autoencoder_rae.py

+        if encoder_hidden_size is None:
+            raise ValueError(f"Encoder '{encoder_cls}' must define `.hidden_size` attribute.")
+
+        decoder_config = SimpleNamespace(


Why do we need this?

Not addressed yet,

sayakpaul · 2026-02-23T03:19:16Z

src/diffusers/models/autoencoders/autoencoder_rae.py

+    - trainable_cls_token
+    """
+
+    def __init__(self, config, num_patches: int):


We should split out the config and expand the __init__ args here. That's how it's done in diffusers.

This is not addressed.

bytetriper · 2026-02-23T03:29:20Z

@bytetriper thanks! What would be the quickest way to validate if the implementation is correct? We can do a quick value assertion test between the original model and the converted model on the same inputs. Would you be able to do it?

I tested and the converted model produce identical output on my end up to some small numerical differences. Just want to make sure it also has the same behavior on other's end:)

I generally agree that we should have encoder in the ckpt as well. Can help for conversion afterewards

sayakpaul · 2026-02-23T04:17:50Z

Cool then. I will give you a heads up when the PR is ready for another look. Thank you!

kashif · 2026-02-23T09:49:32Z

@bytetriper i sent you some fixes to the weights if you can kindly merge

bytetriper · 2026-02-23T16:42:07Z

@kashif Merged!

…aults

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.

sayakpaul · 2026-02-27T10:33:23Z

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

+specific language governing permissions and limitations under the License.
+-->
+
+# AutoencoderRAE


@stevhliu could you check out the docs?

sayakpaul · 2026-02-27T10:34:21Z

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

+    "nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08"
+).to("cuda").eval()
+
+image = Image.open("cat.png").convert("RGB").resize((224, 224))


Can we use an example snippet that just works? In that case, we should load the image directly from a public URL and then use it further. We can leverage load_image for this (from diffusers import load_image).

sayakpaul · 2026-02-27T10:35:01Z

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

+# Latent normalization is handled automatically inside encode/decode
+# when the checkpoint config includes latents_mean/latents_std.


examples/research_projects/autoencoder_rae/README.md

sayakpaul · 2026-02-27T10:37:35Z

src/diffusers/models/autoencoders/autoencoder_rae.py

+        self.model.layernorm.bias = None
+
+
+    def forward(self, images: torch.Tensor, requires_grad: bool = False) -> torch.Tensor:


Should the users need to pass requires_grad or can we just do it based on self.training?

src/diffusers/models/autoencoders/autoencoder_rae.py

sayakpaul · 2026-02-27T10:39:42Z

src/diffusers/models/autoencoders/autoencoder_rae.py

+    - trainable_cls_token
+    """
+
+    def __init__(self, config, num_patches: int):


This is not addressed.

sayakpaul · 2026-02-27T10:40:29Z

src/diffusers/models/autoencoders/autoencoder_rae.py

+        if encoder_hidden_size is None:
+            raise ValueError(f"Encoder '{encoder_cls}' must define `.hidden_size` attribute.")
+
+        decoder_config = SimpleNamespace(


Not addressed yet,

sayakpaul · 2026-02-27T10:41:09Z

tests/models/autoencoders/test_models_autoencoder_rae.py

+_ENCODER_TYPES["tiny_test"] = TinyTestEncoder
+
+
+class AutoencoderRAETests(unittest.TestCase):


Can we not use the existing tester mixin for this?

sayakpaul · 2026-02-27T10:41:44Z

tests/models/autoencoders/test_models_autoencoder_rae.py

+        self.assertTrue(torch.allclose(z_eval_1, z_eval_2, atol=1e-6, rtol=1e-5))
+
+
+@slow


We can skip this test suite for now. Let's first have enough growth for the model.

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

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 23, 2026

View reviewed changes

sayakpaul requested a review from dg845 February 23, 2026 03:25

fix entrypoint for instantiating the AutoencoderRAE

c68b812

kashif added 2 commits February 23, 2026 09:59

added encoder_image_size config

61885f3

undo last change

28a02eb

kashif added 8 commits February 25, 2026 13:38

fixes from pretrained weights

b297868

Merge branch 'main' into rae

7debd07

cleanups

b3ffd63

address reviews

dca5923

Merge branch 'rae' of https://github.com/Ando233/diffusers into rae

c71cb44

fix train script to use pretrained

5c85781

fix conversion script review

d965cab

latebt normalization buffers are now always registered with no-op def…

663b580

…aults

kashif requested a review from sayakpaul February 26, 2026 12:22

sayakpaul reviewed Feb 27, 2026

View reviewed changes

kashif and others added 2 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>


		For latent normalization, use `latents_mean` and `latents_std` (matching other diffusers autoencoders).

		See `examples/research_projects/autoencoder_rae/train_autoencoder_rae.py` for a stage-1 style training script


		`encoder_cls` supports `"dinov2"`, `"siglip2"`, and `"mae"`.

		For latent normalization, use `latents_mean` and `latents_std` (matching other diffusers autoencoders).

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

		# Latent normalization is handled automatically inside encode/decode
		# when the checkpoint config includes latents_mean/latents_std.

		self.model.layernorm.bias = None


		def forward(self, images: torch.Tensor, requires_grad: bool = False) -> torch.Tensor:

		_ENCODER_TYPES["tiny_test"] = TinyTestEncoder


		class AutoencoderRAETests(unittest.TestCase):

		self.assertTrue(torch.allclose(z_eval_1, z_eval_2, atol=1e-6, rtol=1e-5))


		@slow

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!

Ando233 commented Feb 17, 2026

Uh oh!

kashif commented Feb 17, 2026

Uh oh!

bytetriper commented Feb 22, 2026

Uh oh!

sayakpaul commented Feb 23, 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!

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bytetriper commented Feb 23, 2026

Uh oh!

sayakpaul commented Feb 23, 2026

Uh oh!

kashif commented Feb 23, 2026

Uh oh!

bytetriper commented Feb 23, 2026

Uh oh!

HuggingFaceDocBuilderDev commented Feb 27, 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!

Uh oh!

Choose a reason for hiding this comment

Uh oh!