nvidia/C-RADIOv4-1D-H

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 Namespace9 10import torch11from torch import nn12import torch.nn.functional as F13 14from .adaptor_registry import adaptor_registry, dict_t, state_t15 16from .adaptor_generic import GenericAdaptor17from .utils import rank_gate18 19class OpenCLIP_RADIO(GenericAdaptor):20    def __init__(self, main_config: Namespace, adaptor_config: dict_t, state: state_t):21        super().__init__(main_config, adaptor_config, state)22 23        import open_clip24        with rank_gate():25            self.oc_model = open_clip.create_model_from_pretrained(26                model_name=adaptor_config['model'],27                pretrained=adaptor_config['pretrained'],28                return_transform=False,29            )30        # Unload these parameters31        self.oc_model.visual = None32 33        self.tokenizer = open_clip.get_tokenizer(model_name=adaptor_config['model'])34 35    def encode_text(self, text, normalize: bool = False):36        return self.oc_model.encode_text(text, normalize=normalize)37 38 39@adaptor_registry.register_adaptor("open_clip")40def create_open_clip_adaptor(main_config: Namespace, adaptor_config: dict_t, state: state_t):41    return OpenCLIP_RADIO(main_config, adaptor_config, state)42