Update references and fix demo code