Small UI working changes