GenAI-Arena / model /models /imagenhub_models.py
DongfuJiang's picture
update
bd0bce1
raw
history blame
330 Bytes
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)