nvidia/C-RADIO

1# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.2#3# NVIDIA CORPORATION and its licensors retain all intellectual property4# and proprietary rights in and to this software, related documentation5# and any modifications thereto.  Any use, reproduction, disclosure or6# distribution of this software and related documentation without an express7# license agreement from NVIDIA CORPORATION is strictly prohibited.8from argparse import Namespace9from typing import Dict, Any10 11import torch12 13from .adaptor_generic import GenericAdaptor, AdaptorBase14 15dict_t = Dict[str, Any]16state_t = Dict[str, torch.Tensor]17 18 19class AdaptorRegistry:20    def __init__(self):21        self._registry = {}22 23    def register_adaptor(self, name):24        def decorator(factory_function):25            if name in self._registry:26                raise ValueError(f"Model '{name}' already registered")27            self._registry[name] = factory_function28            return factory_function29        return decorator30 31    def create_adaptor(self, name, main_config: Namespace, adaptor_config: dict_t, state: state_t) -> AdaptorBase:32        if name not in self._registry:33            return GenericAdaptor(main_config, adaptor_config, state)34        return self._registry[name](main_config, adaptor_config, state)35 36# Creating an instance of the registry37adaptor_registry = AdaptorRegistry()38