Replaced code with a cleaner version of code for demos