import imagen_hub

class ImagenHubModel():
    def __init__(self, model_name):
        self.model = imagen_hub.load(model_name)
    
    def __call__(self, *args, **kwargs):
        return self.model.infer_one_image(*args, **kwargs)
    
def load_imagenhub_model(model_name, model_type=None):
    return ImagenHubModel(model_name)