Cleaned up code for demos