A from-scratch implementation of a Latent Diffusion Model using Transformers, trained on Minecraft gameplay footage. This project explores the intersection of Variational AutoEncoders (VAE) and Diffusion models for video frame synthesis.
This project implements a complete generative pipeline:
- VAE: A Transformer-based Variational AutoEncoder that compresses 640x352 frames into a compact latent space.
- Latent Diffusion: A timestep-guided Transformer model that learns to denoise latent vectors.
- RoPE: Uses Rotary Positional Embeddings for both the VAE and the Diffuser to handle spatial relationships in the patch-based grid.
- Custom RoPE Attention: 2D Rotary Positional Embeddings implemented for high-resolution grids.
- Latent Training: Training on precomputed VAE latents for significant speedup.
- DDIM Sampling: Optimized sampler for faster inference.
- Input: 640x352 RGB images.
- Patches: 16x16 pixels (total 880 tokens).
- Transformer: 6 layers for VAE, 14 layers for Diffuser (1024 inner dim).
- Positional Encoding: 2D RoPE.
Note
The current results exhibit a "patchy" look. This is a known artifact of the current VAE implementation where non-overlapping convolutions are used for patch extraction and reconstruction.
- Dataset: ~48 hours of Minecraft gameplay footage (20 fps).
- Calcul: Trained on NVIDIA RTX Pro 6000 Blackwell (96GB VRAM).
- Time: 2h for VAE convergence, 48h for Diffusion model.
I know this is very slow, but for a first attempt at training a diffusion model from scratch, I think it's pretty good.
pip install torch torchvision einops tqdm pillow pandaspython train_vae.pyExtract latents to speed up diffusion training:
# See notebook logic for precomputation scriptpython train_diffusion.pypython generate.py- Smoothing Boundaires: Implementing overlapping patches or a convolutional refinement stage in the VAE decoder to eliminate the "patchy" look.
- Temporal Consistency: Expanding the model to handle sequences of frames rather than single images.
Developed by Gabriel Rochet (2026)