Small cleanup on stage model. Hopefully my code sticks now.