toilaluan/raev2-compressor

This is an adapter for RAEv2 checkpoint. This adapter further compresses feature space of RAE with pool size is 2. That means if input is 256x256, original RAE produces 256x1024 latent space. With this adapter, we trained on AutoEncoder objective, produces only 64x1024 latent space, which mostly remove redundancy of original latent space.

from transformers import AutoModel

# Load rae
rae = AutoModel.from_pretrained(
    "toilaluan/raev2-dinov3l-k7",
    trust_remote_code=True,
).eval().to("cuda")

compact_rae = AutoModel.from_pretrained(
    "toilaluan/rae-compressor",
    rae=rae,
    trust_remote_code=True,
).eval().to("cuda")

images = torch.rand(1, 3, 256, 256, device="cuda")
latents = model.encode(images)       # [1, 1024, 8, 8]
reconstructions = model.decode(latents)  # [1, 3, 256, 256]