Project Overview
Implemented a from-scratch Stable Diffusion 3 style pipeline with transformer-based MMDiT to generate 64x64 text-conditioned images.
Problem Definition & Goals
- Problem: Needed to move beyond MNIST-scale diffusion and implement latent diffusion plus text conditioning in a practical training setup.
- Goal 1: Recreate SD3 core architecture with MMDiT and VAE components.
- Goal 2: Build end-to-end text-to-image training and inference at 64x64 resolution.
- Goal 3: Evaluate model quality with objective metrics instead of visual-only inspection.
Key Features & Contributions
- MMDiT Core: Implemented transformer backbone with joint attention for image-text interaction.
- Conditioning Path: Added CLIP-L and adaLN-zero based time/text conditioning.
- Rectified Flow + VAE: Integrated rectified flow training and beta-VAE latent compression (VAE MSE 0.0002).
- Scalable Training Stack: Used LAION/CC3M strategy, Hugging Face streaming, and W&B + RunPod monitoring.
- Evaluation Suite: Added FID, CLIP-FID, IS, CLIPScore, PSNR, and MSE for multidimensional validation.
Technical Challenges & Solutions
- Prompt Overfitting: Small 30k-scale data caused repetitive pattern generation independent of prompts.
- Generalization Fix: Expanded to CC3M-scale data and tuned VAE-to-diffusion data ratio.
- Slow Convergence: Diffusion loss decreased slowly and training became unstable in long runs.
- Stability Fix: Applied logit-normal sampling, min-SNR weighting, and checkpoint-resume workflow.
Results & Learnings
- Result: Delivered a working SD3-style text-conditioned generation model with MMDiT and rectified flow.
- Result: Established a fully quantitative benchmarking pipeline for generation quality and alignment.
- Result: Completed an end-to-end 64x64 text-to-image training/inference stack.
- Learning: Strengthened architecture engineering, data strategy design, and resource-aware training operations.