from typing import List
import gradio as gr
from PIL import Image
from models import load_transformers
from libs.model_list import model_list



def multiple_image_captioning(input_img: Image.Image, question:str) -> List:
    results = []
    for model_name, pretrained_paths in model_list.items():

        for pretrained_path in pretrained_paths:
            try:
                process = load_transformers(name=model_name, model_pretrain=pretrained_path)
                
                if question == '':
                    text = process.image_captioning(input_img)
                else:
                    text = process.visual_question_answering(input_img)
            except Exception as e:
                text = str(e)
                
            results.append(text)
    return results



question_text_Box = gr.Textbox(label="Question")

outputs = [] 
for model_name, pretrained_paths in model_list.items():
        for pretrained_path in pretrained_paths:
            outputs.append(gr.Textbox(label=model_name, info=pretrained_path))
            
            
demo = gr.Interface(fn=multiple_image_captioning, 
                    inputs=[gr.Image(type='pil'), question_text_Box],
                    outputs=outputs,
                )

demo.launch()