Scale-RAE provides official implementations for scaling text-to-image diffusion transformers using Representation Autoencoders. It enables efficient large-scale pretraining and finetuning, targeting researchers and engineers developing advanced generative AI models. The project offers GPU inference and robust TPU training capabilities, facilitating the creation of high-fidelity image generation systems.

How It Works

The core innovation lies in a two-stage training methodology. Stage 1 involves large-scale pretraining utilizing pre-existing Large Language Models (LLMs) and randomly initialized Diffusion Transformers (DiTs). Stage 2 refines these models through finetuning on curated instruction datasets. This approach leverages Representation Autoencoders to enhance the scaling efficiency and performance of DiTs in text-to-image generation, supported by decoders like SigLIP-2 and WebSSL.

Quick Start & Requirements

Installation involves cloning the repository, creating a Python 3.10 conda environment, and installing the package. GPU is required for inference, with a command-line interface provided. Large-scale training is optimized for TPUs, with a dedicated setup guide available. Models are automatically downloaded from HuggingFace.

• Install: pip install -e .
• Inference: python inference/cli.py t2i --prompt "..."
• Prerequisites: Python 3.10, GPU (inference), TPU (training).
• Docs: Inference Guide, Training Guide, TPU Setup Guide.

Highlighted Details

• Offers pre-trained models combining Qwen2.5 LLMs with DiT models up to 9.8B parameters, using SigLIP-2 or WebSSL decoders.
• Supports both efficient GPU inference and scalable TPU training via SPMD/FSDP.
• Two-stage training pipeline for robust model development.
• Models are readily available via HuggingFace repositories.

Maintenance & Community

The project originates from New York University, with core contributors listed among the paper's authors. No specific community channels (e.g., Discord, Slack) or active maintenance signals beyond the initial release are detailed in the README.

Licensing & Compatibility

Released under the MIT License, permitting broad use, modification, and distribution, including for commercial applications.

Limitations & Caveats

The primary focus is on GPU inference and TPU training, suggesting limited optimization for CPU-based inference. Setting up TPUs for large-scale training may require significant configuration effort as detailed in the dedicated guide. The README does not detail specific performance benchmarks or known bugs.