Refactor and clean diagram rendering

Some diagrams were not rendering properly, and there were stray buttons. This merge fixes the issue as well as refactor some code