PCGao/MatchAnything

1import sys2from abc import ABCMeta, abstractmethod3from torch import nn4from copy import copy5import inspect6from huggingface_hub import hf_hub_download7 8 9class BaseModel(nn.Module, metaclass=ABCMeta):10    default_conf = {}11    required_inputs = []12 13    def __init__(self, conf):14        """Perform some logic and call the _init method of the child model."""15        super().__init__()16        self.conf = conf = {**self.default_conf, **conf}17        self.required_inputs = copy(self.required_inputs)18        self._init(conf)19        sys.stdout.flush()20 21    def forward(self, data):22        """Check the data and call the _forward method of the child model."""23        for key in self.required_inputs:24            assert key in data, "Missing key {} in data".format(key)25        return self._forward(data)26 27    @abstractmethod28    def _init(self, conf):29        """To be implemented by the child class."""30        raise NotImplementedError31 32    @abstractmethod33    def _forward(self, data):34        """To be implemented by the child class."""35        raise NotImplementedError36 37    def _download_model(self, repo_id=None, filename=None, **kwargs):38        """Download model from hf hub and return the path."""39        return hf_hub_download(40            repo_type="model",41            repo_id=repo_id,42            filename=filename,43        )44 45 46def dynamic_load(root, model):47    module_path = f"{root.__name__}.{model}"48    module = __import__(module_path, fromlist=[""])49    classes = inspect.getmembers(module, inspect.isclass)50    # Filter classes defined in the module51    classes = [c for c in classes if c[1].__module__ == module_path]52    # Filter classes inherited from BaseModel53    classes = [c for c in classes if issubclass(c[1], BaseModel)]54    assert len(classes) == 1, classes55    return classes[0][1]56    # return getattr(module, 'Model')57