| # CIFAR-10 Diffusion Model | |
| A lightweight diffusion model trained from scratch on the CIFAR-10 dataset in just 14.5 minutes using PyTorch. | |
| ## Model Description | |
| This is a **SimpleUNet-based diffusion model** trained to generate 32x32 RGB images similar to the CIFAR-10 dataset. The model demonstrates the fundamentals of diffusion-based image generation with a compact architecture suitable for educational purposes and quick experimentation. | |
| ### Key Features | |
| - 🚀 **Fast Training**: Complete training in under 15 minutes on RTX 3060 | |
| - 💾 **Lightweight**: Only 16.8M parameters (~64MB model size) | |
| - 🎯 **Educational**: Clean, well-documented code for learning diffusion models | |
| - ⚡ **Efficient Inference**: Generate images in seconds on consumer GPUs | |
| ## Model Details | |
| | Attribute | Value | | |
| |-----------|-------| | |
| | **Architecture** | SimpleUNet with ResNet blocks + Attention | | |
| | **Parameters** | 16,808,835 | | |
| | **Dataset** | CIFAR-10 (50,000 training images) | | |
| | **Image Size** | 32×32 RGB | | |
| | **Training Steps** | 7,820 (20 epochs × 391 batches) | | |
| | **Training Time** | 14.54 minutes | | |
| | **Hardware** | NVIDIA RTX 3060 (0.43GB VRAM used) | | |
| | **Framework** | PyTorch 2.0+ | | |
| ## Quick Start | |
| ### Installation | |
| ```bash | |
| pip install torch torchvision matplotlib tqdm pillow numpy | |
| ``` | |
| ### Basic Usage | |
| ```python | |
| import torch | |
| import matplotlib.pyplot as plt | |
| # Load model | |
| checkpoint = torch.load('complete_diffusion_model.pth') | |
| model = SimpleUNet(**checkpoint['model_config']) | |
| model.load_state_dict(checkpoint['model_state_dict']) | |
| model.eval() | |
| # Initialize scheduler | |
| scheduler = DDPMScheduler(**checkpoint['diffusion_config']) | |
| # Generate images | |
| @torch.no_grad() | |
| def generate_images(model, scheduler, num_images=4): | |
| device = next(model.parameters()).device | |
| images = torch.randn(num_images, 3, 32, 32).to(device) | |
| for t in range(999, -1, -20): # 50 denoising steps | |
| timestep = torch.full((num_images,), t, device=device) | |
| noise_pred = model(images, timestep) | |
| # Simplified DDPM step | |
| alpha_t = scheduler.alpha_cumprod[t] | |
| alpha_prev = scheduler.alpha_cumprod[t-20] if t >= 20 else 1.0 | |
| pred_x0 = (images - torch.sqrt(1-alpha_t) * noise_pred) / torch.sqrt(alpha_t) | |
| images = torch.sqrt(alpha_prev) * pred_x0 + torch.sqrt(1-alpha_prev) * noise_pred | |
| return images | |
| # Generate and display | |
| generated = generate_images(model, scheduler) | |
| ``` | |
| ## Training Details | |
| - **Loss Function**: MSE between predicted and actual noise | |
| - **Optimizer**: AdamW (lr=1e-4, weight_decay=1e-6) | |
| - **Scheduler**: CosineAnnealingLR | |
| - **Batch Size**: 128 | |
| - **Final Loss**: 0.0363 (73% reduction from initial) | |
| - **Diffusion Steps**: 1000 (linear beta schedule) | |
| ## Performance | |
| ### Training Loss Curve | |
| The model shows excellent convergence: | |
| - **Epoch 1**: 0.1349 → **Epoch 20**: 0.0363 | |
| - **Best Loss**: 0.0358 (Epoch 19) | |
| - **Stable convergence** without overfitting | |
| ### Generation Quality | |
| - ✅ Captures CIFAR-10 color distributions | |
| - ✅ Generates diverse, non-repetitive outputs | |
| - ⚠️ Abstract patterns (needs longer training for object recognition) | |
| - 🎯 Suitable for color/texture generation tasks | |
| ## Files in this Repository | |
| | File | Description | Size | | |
| |------|-------------|------| | |
| | `complete_diffusion_model.pth` | Full model with config and weights | ~64MB | | |
| | `diffusion_model_final.pth` | Training checkpoint (epoch 20) | ~64MB | | |
| | `model_info.json` | Training metadata and hyperparameters | <1KB | | |
| | `inference_example.py` | Complete inference script with model classes | ~5KB | | |
| ## Model Architecture | |
| ``` | |
| SimpleUNet( | |
| time_embedding: TimeEmbedding(128) | |
| encoder: 3 ResNet blocks with downsampling | |
| middle: ResNet + Self-Attention + ResNet | |
| decoder: 3 ResNet blocks with upsampling | |
| output: GroupNorm → SiLU → Conv2d | |
| ) | |
| ``` | |
| ## Use Cases | |
| - 🎓 **Educational**: Learn diffusion model fundamentals | |
| - 🔬 **Research**: Baseline for diffusion experiments | |
| - 🎨 **Art**: Generate abstract textures and patterns | |
| - ⚡ **Prototyping**: Quick diffusion model testing | |
| ## Limitations & Improvements | |
| ### Current Limitations | |
| - Generates abstract patterns rather than recognizable objects | |
| - Trained on small 32×32 resolution | |
| - Limited to 20 training epochs | |
| ### Suggested Improvements | |
| 1. **Extended Training**: 50-100 epochs for better object generation | |
| 2. **Larger Architecture**: Increase model capacity | |
| 3. **Advanced Sampling**: Implement DDIM or DPM-Solver++ | |
| 4. **Higher Resolution**: Train on 64×64 or 128×128 images | |
| 5. **Better Datasets**: Use CelebA-HQ or custom datasets | |
| ## Citation | |
| ```bibtex | |
| @misc{cifar10-diffusion-2025, | |
| title={CIFAR-10 Diffusion Model: Fast Training Implementation}, | |
| author={Karthik}, | |
| year={2025}, | |
| publisher={Hugging Face}, | |
| howpublished={\url{https://huggingface.co/karthik-2905/DiffusionPretrained}} | |
| } | |
| ``` | |
| ## License | |
| MIT License - Free for research and commercial use. | |
| --- | |
| **🚀 Want to train your own?** Check out the [full implementation](https://github.com/GruheshKurra/DiffusionModelPretrained) with Jupyter notebooks and step-by-step training code! | |
| **📊 Training Stats**: 16.8M params • 14.5min training • RTX 3060 • PyTorch 2.0 |