A high-performance, memory-optimized image generation interface built specifically for Google Colab environments. This project provides a user-friendly Gradio interface for generating high-quality images using Hugging Face's Diffusers library with advanced memory management and performance optimizations.
- Aggressive cleanup: Automatic GPU memory clearing between operations
- CPU offloading: Intelligent memory offloading to prevent OOM errors
- Sequential offloading: Maximum memory efficiency for large models
- Real-time monitoring: Live GPU memory usage tracking
- TF32 acceleration: Faster computation on modern GPUs
- xFormers attention: Memory-efficient attention mechanisms
- VAE optimizations: Slicing and tiling for large images
- CUDA optimizations: Benchmark mode and cache management
- Modern interface: Clean, professional Gradio UI
- Batch generation: Generate multiple images at once
- Memory monitor: Real-time system resource tracking
- Download support: Easy image saving and sharing
- Error recovery: Graceful handling of memory issues
- Multiple offloading strategies: Balance between speed and memory
- Session management: Better handling of runtime limitations
- Automatic device detection: Works on GPU and CPU
- Share links: Generate public URLs for sharing
-
Enable GPU Runtime
Runtime → Change runtime type → T4 GPU → Save -
Install Dependencies
!pip install torch !pip install -U xformers --index-url https://download.pytorch.org/whl/cu126 !pip install diffusers==0.32.2 !pip install transformers==4.49
-
Run the Application
# Download and run the optimized script !wget https://github.com/abhi-0203/Optimized-AI-Image-Generator-for-Google-Colab/blob/main/optimized_genai_colab.ipynb %run optimized_genai_colab.py
- Load a model: Enter a Hugging Face model ID (e.g.,
runwayml/stable-diffusion-v1-5) - Configure settings:
- Enable CPU offloading if you encounter memory issues
- Adjust image dimensions and inference steps
- Enter prompts: Write your image description and negative prompts
- Generate: Click generate and wait for your image!
runwayml/stable-diffusion-v1-5 # Lightweight and fast
stabilityai/stable-diffusion-2-1-base # Good quality/speed balance
dreamlike-art/dreamlike-diffusion-1.0 # Artistic style
stabilityai/stable-diffusion-xl-base-1.0 # High quality
stabilityai/sdxl-turbo # Fast SDXL variant
runwayml/stable-diffusion-v1-5 # Still great for speed
# Level 1: GPU only (fastest, highest memory usage)
generator.load_model("model-id", use_cpu_offload=False, use_sequential_offload=False)
# Level 2: CPU offloading (balanced)
generator.load_model("model-id", use_cpu_offload=True, use_sequential_offload=False)
# Level 3: Sequential offloading (slowest, lowest memory usage)
generator.load_model("model-id", use_cpu_offload=False, use_sequential_offload=True)| Parameter | Recommended Range | Description |
|---|---|---|
| Steps | 15-30 | Quality vs speed trade-off |
| Guidance Scale | 7.0-12.0 | Prompt adherence strength |
| Image Size | 512-1024px | Multiples of 64 work best |
| Batch Size | 1-4 | Number of images to generate |
- Python: 3.10+
- PyTorch: 2.0+
- CUDA: 11.8+ (for GPU acceleration)
- RAM: 8GB+ system RAM
- VRAM: 6GB+ for most models
diffusers>=0.21.0
transformers>=4.25.0
accelerate>=0.16.0
xformers>=0.0.16
gradio>=4.0.0
torch>=2.0.0
The application is built with:
- Core:
OptimizedImageGeneratorclass with smart memory management - Interface: Gradio web interface with custom CSS
- Backend: Hugging Face Diffusers pipeline
- Optimization: xFormers, CPU offloading, VAE optimizations
| Configuration | Model | Image Size | Steps | Time (T4) | Memory Usage |
|---|---|---|---|---|---|
| Standard | SD 1.5 | 512x512 | 20 | ~8s | ~4GB |
| Optimized | SD 1.5 | 512x512 | 20 | ~6s | ~2.5GB |
| CPU Offload | SDXL | 1024x1024 | 25 | ~45s | ~8GB |
| Sequential | SDXL | 1024x1024 | 25 | ~60s | ~4GB |
# Solutions (in order of preference):
1. Enable CPU offloading
2. Reduce image dimensions (1024→768→512)
3. Lower inference steps (30→20→15)
4. Use sequential offloading
5. Restart runtime and try again# Common fixes:
1. Check internet connection
2. Verify model ID spelling
3. Clear memory before loading new model
4. Try alternative model variants# Optimization steps:
1. Ensure GPU runtime is enabled
2. Verify CUDA availability
3. Reduce image size and steps
4. Use CPU offloading only when necessary- Monitor usage: Use the built-in memory monitor
- Clear regularly: Click "Clear Memory" between model switches
- Start small: Begin with 768x768 images
- Use offloading: Enable for models >6GB
Contributions are welcome! Please feel free to submit a Pull Request. For major changes, please open an issue first to discuss what you would like to change.
- Fork the repository
- Create a feature branch:
git checkout -b feature/amazing-feature - Make your changes
- Test in Colab: Ensure compatibility
- Commit changes:
git commit -m 'Add amazing feature' - Push to branch:
git push origin feature/amazing-feature - Open a Pull Request
- Follow PEP 8 guidelines
- Use type hints where appropriate
- Add docstrings for new functions
- Test changes in Google Colab
- Hugging Face: For the amazing Diffusers library
- Stability AI: For Stable Diffusion models
- Google Colab: For providing free GPU access
- Gradio: For the excellent web interface framework
- Issues: GitHub Issues
- Support for ControlNet models
- Inpainting and outpainting features
- Image-to-image generation
- Advanced prompt engineering tools
- Multi-model comparison interface
- API endpoint for programmatic access
Star ⭐ this repository if you find it helpful!
Made with ❤️ for the AI community