import gradio as gr
from llava.mm_utils import get_model_name_from_path
from llava.model.builder import load_pretrained_model
from llava.eval.run_llava import eval_model

model_path = "liuhaotian/llava-v1.5-7b"

model_name = get_model_name_from_path(model_path)
tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name, load_4bit=True)

def predict(input_img, prompt):

    args = type('Args', (), {
        "model_path": model_path,
        "model_base": None,
        "model_name": model_name,
        "query": prompt,
        "conv_mode": None,
        "image_file": input_img,
        "sep": ",",
        "temperature": 0.2,
        "top_p": None,
        "num_beams": 1,
        "max_new_tokens": 512
    })()

    return eval_model(args, tokenizer, model, image_processor, context_len)

gradio_app = gr.Interface(
    fn=predict,
    inputs=[gr.Image(type="filepath"), "text"],
    outputs="text"
)

gradio_app.launch()