Add __init__.py back in