Cleaned up main.py a bit.