---

name: vae-training description: "Train and inspect DALLE-pytorch discrete VAE workflows, VAE checkpoints, and pretrained VAE alternatives." metadata: disco-role: operating disable-model-invocation: true license: MIT

VAE training

Use this sub-skill when the user asks about DiscreteVAE, VAE codebooks, VAE image folders, VAE checkpoint files, OpenAI's pretrained VAE, Taming Transformers VQGAN, or the VAE stage before DALL-E transformer training.

Read first

• references/workflows.md for VAE API recipes, command construction, checkpoint contents, and pretrained VAE choices.
• references/troubleshooting.md for image-size assertions, OpenAI VAE torch limits, W&B/GPU side effects, and VQGAN path errors.
• scripts/tiny_vae_api_smoke.py when you need a no-download CPU check.
• scripts/build_train_vae_command.py when a user needs a reproducible VAE training command template.

Typical routes

Request	Action
"Train the VAE" or "prepare vae.pt"	Validate image folder shape, build a command with scripts/build_train_vae_command.py, and warn about CUDA/W&B/checkpoint side effects before running long training.
"Use DiscreteVAE in Python"	Use the API recipe in references/workflows.md; verify with scripts/tiny_vae_api_smoke.py.
"OpenAI pretrained VAE fails"	Check torch version; source asserts torch <=1.10. Prefer a trained DiscreteVAE or VQGanVAE unless the user accepts a legacy environment.
"Use VQGAN VAE"	Require both model checkpoint and YAML config paths when not using the default download path; route backend/download questions to root troubleshooting.
"What is in vae.pt?"	Explain hparams and weights, and that DALL-E training reconstructs DiscreteVAE(**hparams) before loading weights.

Minimal API pattern

import torch
from dalle_pytorch import DiscreteVAE

vae = DiscreteVAE(
    image_size=256,
    num_layers=3,
    num_tokens=8192,
    codebook_dim=512,
    hidden_dim=64,
    num_resnet_blocks=1,
)
images = torch.randn(4, 3, 256, 256)
loss = vae(images, return_loss=True)
loss.backward()
indices = vae.get_codebook_indices(images)
recon = vae.decode(indices)

Keep detailed constructor notes and training command options in references/workflows.md, not in this router.

Boundary notes

• DALL-E transformer training after VAE preparation belongs to ../dalle-training/SKILL.md.
• Image generation and CLIP ranking belong to ../generation-and-ranking/SKILL.md.
• DeepSpeed/Apex/Horovod/Docker/CUDA setup belongs to ../distributed-and-backends/SKILL.md.
• Do not instantiate pretrained VAE wrappers merely to inspect the package; they can download model files or assert on torch version.

Validation checklist

Before approving a VAE workflow answer:

• image size is a power of two and matches the intended downstream DALL-E VAE;
• image mode (RGB vs transparent RGBA) matches channels;
• checkpoint side effects and output filenames are explicit;
• W&B/account/network behavior is acknowledged;
• if a full script command is produced, the user has a checkout or script copy that can run it;
• package-level API alternatives are offered for pip-only users.