EdgeStyleGAN is a lightweight variant of StyleGAN optimized for efficient deployment on resource-constrained devices such as mobile phones and edge computing devices. Our approach achieves significant reductions in model size and inference time while preserving high-quality image generation capabilities.
- Lightweight Architecture: Up to X% reduction in parameters compared to original StyleGAN
- Fast Inference: Y× faster inference speed for real-time applications
- Quality Preservation: Maintains perceptual quality as measured by FID scores and human evaluation
- Mobile-Optimized: Designed for deployment on mobile and edge devices
- Real-Time Generation: Enables real-time image synthesis applications
Our optimization strategy combines three key techniques:
- Structured Pruning: Systematic removal of redundant network components
- Weight Quantization: Reduced precision representation of model weights
- Knowledge Distillation: Transfer learning from full-size StyleGAN teacher model
- Personalization: Real-time avatar and character generation
- Augmented Reality: Live image synthesis for AR applications
- Creative Tools: Mobile apps for artistic content creation
- Gaming: On-device asset generation for games
- Social Media: Real-time filters and effects
Python 3.8+
PyTorch 1.9+
CUDA 10.2+ (for GPU acceleration)# Clone the repository
git clone https://github.com/jash0803/edgestylegan.git
cd edgestylegan
# Create virtual environment
python -m venv venv
source venv/bin/activate # On Windows: venv\Scripts\activate
# Install dependencies
pip install -r requirements.txtimport torch
from edgestylegan import EdgeStyleGAN
# Load pre-trained model
model = EdgeStyleGAN.from_pretrained('edgestylegan-ffhq-256')
# Generate random images
with torch.no_grad():
latent = torch.randn(1, 512) # Random latent vector
image = model.generate(latent)
# Save generated image
from torchvision.utils import save_image
save_image(image, 'generated.png', normalize=True, range=(-1, 1))# Initialize for real-time inference
model.eval()
model = model.cuda() # Move to GPU if available
# Generate multiple images quickly
batch_size = 4
latents = torch.randn(batch_size, 512).cuda()
start_time = time.time()
with torch.no_grad():
images = model.generate(latents)
end_time = time.time()
print(f"Generated {batch_size} images in {end_time - start_time:.3f}s")| Model | Dataset | Resolution | Parameters | FID ↓ | Inference Time | Download |
|---|---|---|---|---|---|---|
| EdgeStyleGAN-FFHQ | FFHQ | 256×256 | 12.5M | 8.2 | 15ms | Link |
| EdgeStyleGAN-FFHQ | FFHQ | 512×512 | 25.1M | 6.8 | 32ms | Link |
| EdgeStyleGAN-LSUN | LSUN-Car | 256×256 | 12.8M | 12.1 | 16ms | Link |
# Download and prepare FFHQ dataset
python scripts/prepare_data.py --dataset ffhq --resolution 256
# For custom datasets
python scripts/prepare_data.py --dataset custom --data_path /path/to/images --resolution 256# Train from scratch
python train.py --config configs/edgestylegan_ffhq_256.yaml
# Train with knowledge distillation
python train.py --config configs/edgestylegan_kd_ffhq_256.yaml --teacher_model stylegan2_ffhq
# Resume training
python train.py --config configs/edgestylegan_ffhq_256.yaml --resume checkpoints/latest.pthKey training parameters in configs/edgestylegan_ffhq_256.yaml:
model:
resolution: 256
latent_dim: 512
pruning_ratio: 0.6
quantization_bits: 8
training:
batch_size: 16
learning_rate: 0.0001
epochs: 1000
optimization:
enable_pruning: true
enable_quantization: true
enable_distillation: true
teacher_weight: 0.5# Calculate FID score
python evaluate.py --model edgestylegan-ffhq-256 --metric fid --real_data_path data/ffhq
# Calculate LPIPS score
python evaluate.py --model edgestylegan-ffhq-256 --metric lpips
# Generate evaluation report
python evaluate.py --model edgestylegan-ffhq-256 --full_evaluation# Benchmark inference speed
python benchmark.py --model edgestylegan-ffhq-256 --device gpu --batch_sizes 1,4,8,16
# Profile memory usage
python benchmark.py --model edgestylegan-ffhq-256 --profile_memory# Convert model to ONNX format
python export_onnx.py --model edgestylegan-ffhq-256 --output edgestylegan.onnx
# Optimize ONNX model for mobile
python optimize_onnx.py --input edgestylegan.onnx --output edgestylegan_mobile.onnx# Convert to TensorFlow Lite
python convert_tflite.py --model edgestylegan-ffhq-256 --output edgestylegan.tflite
# Quantize for mobile deployment
python convert_tflite.py --model edgestylegan-ffhq-256 --output edgestylegan_int8.tflite --quantize| Model | Parameters | MACs | FID Score | Inference Time (ms) |
|---|---|---|---|---|
| StyleGAN2 | 30.0M | 45.2G | 3.2 | 124 |
| EdgeStyleGAN | 12.5M | 18.1G | 8.2 | 15 |
| Reduction | 58.3% | 59.9% | -5.0 | 8.3× |
High-quality face generation results from EdgeStyleGAN compared to StyleGAN2
We welcome contributions! Please see CONTRIBUTING.md for guidelines.
# Install development dependencies
pip install -r requirements-dev.txt
# Run tests
python -m pytest tests/
# Run linting
flake8 edgestylegan/
black edgestylegan/If you use EdgeStyleGAN in your research, please cite our paper:
@article{edgestylegan2024,
title={EdgeStyleGAN: Lightweight and Efficient StyleGAN for Real-Time Image Synthesis on Resource-Constrained Devices},
author={Jash Shah},
journal={Conference/Journal Name},
year={2024}
}This project is licensed under the MIT License - see the LICENSE file for details.
- Original StyleGAN and StyleGAN2 authors for the foundational work
- PyTorch team for the deep learning framework
- Contributors and beta testers
- Primary Author: [email protected]
- Project Issues: GitHub Issues
- Discussions: GitHub Discussions
⭐ Star this repository if you find it useful!