from models.base import Model
import gradio as gr

model = None
predict = None

def load_model(model_name):
    
    global model
    model = Model(model_name)
    global predict
    predict = model.predict

    return model.getModelState()

def fn_gen_text_to_text(x):
    text = predict
    return text

def fn_gen_text_to_image(x):
    image = predict(x)
    return image

def fn_gen_image_to_text(x):
    text = predict(x)
    return text