diffusers/pr_13231/en/api/pipelines/rae_dit.md

RAE DiT

Diffusion Transformers with Representation Autoencoders introduces a two-stage recipe: first train a representation autoencoder (RAE), then train a diffusion transformer on the resulting latent space.

RAEDiTPipeline implements the Stage-2 class-conditional generator in Diffusers. It combines:

• RAEDiT2DModel for latent denoising
• FlowMatchEulerDiscreteScheduler for the denoising trajectory
• AutoencoderRAE for decoding latent samples to RGB images

RAEDiTPipeline expects a Stage-2 checkpoint converted to Diffusers format together with a compatible AutoencoderRAE checkpoint.

This pipeline implements the ImageNet class-conditioned RAE DiT path with plain classifier-free guidance. Upstream AutoGuidance and follow-up internal-guidance variants are out of scope for this pipeline.

Loading a converted pipeline

import torch
from diffusers import RAEDiTPipeline

pipe = RAEDiTPipeline.from_pretrained(
    "path/to/converted-rae-dit-imagenet256",
    torch_dtype=torch.bfloat16,
).to("cuda")

image = pipe(class_labels=[207], num_inference_steps=25).images[0]
image.save("golden_retriever.png")

If the converted pipeline includes an id2label mapping, you can also look up class ids by name:

class_id = pipe.get_label_ids("golden retriever")[0]
image = pipe(class_labels=[class_id], num_inference_steps=25).images[0]

RAEDiTPipeline[[diffusers.RAEDiTPipeline]]

diffusers.RAEDiTPipeline[[diffusers.RAEDiTPipeline]]

Source

Pipeline for class-conditioned image generation in RAE latent space.

__call__diffusers.RAEDiTPipeline.__call__https://github.com/huggingface/diffusers/blob/vr_13231/src/diffusers/pipelines/rae_dit/pipeline_rae_dit.py#L135[{"name": "class_labels", "val": ": int | list[int] | torch.Tensor"}, {"name": "guidance_scale", "val": ": float = 1.0"}, {"name": "guidance_start", "val": ": float = 0.0"}, {"name": "guidance_end", "val": ": float = 1.0"}, {"name": "num_images_per_prompt", "val": ": int = 1"}, {"name": "generator", "val": ": torch.Generator | list[torch.Generator] | None = None"}, {"name": "latents", "val": ": torch.Tensor | None = None"}, {"name": "num_inference_steps", "val": ": int = 50"}, {"name": "output_type", "val": ": str = 'pil'"}, {"name": "return_dict", "val": ": bool = True"}]- class_labels (int, list[int], or torch.Tensor) -- The class ids for the images to generate.

• guidance_scale (float, optional, defaults to 1.0) -- Classifier-free guidance scale. Guidance is enabled when guidance_scale > 1.
• guidance_start (float, optional, defaults to 0.0) -- Lower bound of the normalized timestep interval in which classifier-free guidance is active.
• guidance_end (float, optional, defaults to 1.0) -- Upper bound of the normalized timestep interval in which classifier-free guidance is active.
• num_images_per_prompt (int, optional, defaults to 1) -- Number of images to generate per class label.
• generator (torch.Generator or list[torch.Generator], optional) -- Random generator used for latent sampling.
• latents (torch.Tensor, optional) -- Pre-generated latent noise tensor of shape (batch, channels, height, width).
• num_inference_steps (int, optional, defaults to 50) -- Number of denoising steps.
• output_type (str, optional, defaults to "pil") -- Output format. Choose from "pil", "np", "pt", or "latent".
• return_dict (bool, optional, defaults to True) -- Whether to return an RAEDiTPipelineOutput instead of a tuple.0

The call function to the pipeline for generation.

Parameters:

transformer (RAEDiT2DModel) : Class-conditioned latent transformer used for Stage-2 denoising in RAE latent space.

vae (AutoencoderRAE) : Representation autoencoder used to decode latent samples back to RGB images.

scheduler (FlowMatchEulerDiscreteScheduler) : Flow-matching scheduler used to integrate the latent denoising trajectory.

get_label_ids[[diffusers.RAEDiTPipeline.get_label_ids]]

Source

Map ImageNet-style label strings to class ids.

RAEDiTPipelineOutput[[diffusers.RAEDiTPipelineOutput]]

diffusers.RAEDiTPipelineOutput[[diffusers.RAEDiTPipelineOutput]]

Source

Output class for RAE DiT image generation pipelines.

Parameters:

images (list[PIL.Image.Image] or np.ndarray or torch.Tensor) : Denoised images as PIL images, a NumPy array of shape (batch_size, height, width, num_channels), or a PyTorch tensor of shape (batch_size, num_channels, height, width). Torch tensors may also represent latent outputs when output_type="latent".