extensions_built_in/flex2/__init__.py

from .flex2 import Flex2

AI_TOOLKIT_MODELS = [
    # put a list of models here
    Flex2
]