Update code for demo