Update code from demo