Now works with better model data