Updated demo code