__init__.py

import sys
from inspect import isclass

from .model import Model, model_classes
from .nerf import NeRF
from .fs_nerf import FsNeRF

__all__ = ["Model", "NeRF", "FsNeRF"]


# Register all model classes
for item in __all__:
    var = getattr(sys.modules[__name__], item)
    if isclass(var) and issubclass(var, Model):
        model_classes[item] = var