https://github.com/datacte/sdxl-training-improvements
📊 Research-focused SDXL training framework exploring novel optimization approaches. Goals include enhanced image quality, training stability & comprehensive monitoring. ⭐ Performance-focused research framework.
https://github.com/datacte/sdxl-training-improvements
deep-learning diffusion machine-learning stable-diffusion stable-diffusion-xl
Last synced: about 2 months ago
JSON representation
📊 Research-focused SDXL training framework exploring novel optimization approaches. Goals include enhanced image quality, training stability & comprehensive monitoring. ⭐ Performance-focused research framework.
- Host: GitHub
- URL: https://github.com/datacte/sdxl-training-improvements
- Owner: DataCTE
- License: mit
- Created: 2024-12-19T19:22:37.000Z (10 months ago)
- Default Branch: master
- Last Pushed: 2025-01-13T15:29:26.000Z (9 months ago)
- Last Synced: 2025-04-08T12:46:41.989Z (6 months ago)
- Topics: deep-learning, diffusion, machine-learning, stable-diffusion, stable-diffusion-xl
- Language: Python
- Homepage:
- Size: 1.84 MB
- Stars: 19
- Watchers: 1
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- License: LICENSE
- Authors: AUTHORS
Awesome Lists containing this project
README
# SDXL Training Framework with Novel Research Methods
[](https://www.python.org/downloads/)
[](https://opensource.org/licenses/MIT)
A research-focused SDXL training framework implementing cutting-edge advances in diffusion model training, with emphasis on image quality and training stability.
## Adding New Training Methods
The framework provides a template for easily implementing new training methods:
```python
# src/training/trainers/methods/example_method.py
class ExampleMethodTrainer(SDXLTrainer):
name = "example_method" # Your method's name
def compute_loss(self, batch: Dict[str, Tensor]) -> Dict[str, Any]:
"""Implement your method's loss computation here."""
raise NotImplementedError()
```
To add a new method:
1. Copy `example_method.py` template
2. Implement the `compute_loss()` method with your training logic
3. Register method in config.yaml:
```yaml
training:
method: "your_method_name"
```
The template handles all boilerplate code including:
- Memory optimizations
- Mixed precision training
- Gradient accumulation
- Progress tracking
- Metric logging
## Training Methods
### Flow Matching with Logit-Normal Sampling [[4]](#references)
Advanced training method that eliminates noise scheduling:
```python
# Configure Flow Matching training
training:
method: "flow_matching"
batch_size: 4
learning_rate: 1.0e-6
```
Key benefits:
- 30% faster convergence via optimal transport paths
- Direct velocity field learning reduces instability
- No noise schedule dependencies
- Logit-normal time sampling for better coverage
### NovelAI V3 UNet Architecture [[7]](#references)
State-of-the-art model improvements:
```python
# Enable NovelAI V3 features
training:
prediction_type: "v_prediction"
zero_terminal_snr: true
sigma_max: 20000.0
```
Improvements:
- Zero Terminal SNR training
- Infinite noise approximation (σ_max ≈ 20000)
- Better artifact handling
- Enhanced detail preservation
- V-prediction parameterization
- Stable high-frequency gradients
- Reduced color shifting
- Dynamic Karras schedule
- Adaptive noise spacing
- Improved texture quality
## Training Monitoring
Training progress is monitored using Weights & Biases (wandb):
- Real-time loss tracking
- Generated sample visualization
- Hyperparameter logging
- Custom metric tracking
- Experiment comparison
Enable monitoring in config.yaml:
```yaml
training:
use_wandb: true
```
## Image Quality Improvements
- Enhanced Detail Preservation
- Fine-grained texture generation
- Improved handling of complex patterns
- Better preservation of small objects and features
- Color and Lighting
- More accurate color reproduction
- Enhanced dynamic range in highlights and shadows
- Better handling of complex lighting conditions
- Composition and Structure
- Improved spatial coherence
- Better handling of perspective and depth
- More consistent object proportions
## Requirements
| Component | Version |
|-----------|---------|
| Python | 3.8+ |
| CUDA | 11.7+ |
| VRAM | 24GB+ |
## Installation
```bash
# Clone repository
git clone https://github.com/DataCTE/SDXL-Training-Improvements.git
cd SDXL-Training-Improvements
# Install in development mode with all extras
pip install -e ".[dev,docs]"
# Verify installation
python -c "import src; print(src.__version__)"
```
## Configuration
The training framework is configured through a YAML file. Key configuration sections:
```yaml
# Model configuration
model:
pretrained_model_name: "stabilityai/stable-diffusion-xl-base-1.0"
num_timesteps: 1000
sigma_min: 0.002
sigma_max: 80.0
# Training parameters
training:
batch_size: 4
learning_rate: 4.0e-7
method: "ddpm" # or "flow_matching"
zero_terminal_snr: true
# Dataset configuration
data:
train_data_dir:
- "path/to/dataset1"
- "path/to/dataset2"
```
See [config.yaml](src/config.yaml) for full configuration options.
## Usage Examples
Basic training:
```bash
# Train with default config
python src/main.py --config config.yaml
# Train with custom config
python src/main.py --config my_config.yaml
# Distributed training
torchrun --nproc_per_node=2 src/main.py --config config.yaml
```
## References
4. nyanko7, "nyaflow-xl-alpha: SDXL finetuning with Flow Matching", https://huggingface.co/nyanko7/nyaflow-xl-alpha, 2024
7. Ossa et al., "Improvements to SDXL in NovelAI Diffusion V3", arXiv:2409.15997, 2024
## License
MIT License - see [LICENSE](LICENSE) file.