from panda_gpt import PandaGPT
import gradio as gr

from huggingface_hub import hf_hub_download

vicuna_path = hf_hub_download(repo_id="ningshanwutuobang/ggml-pandagpt-vicuna-merge", filename="ggml-pandagpt-vicuna-q4_1.bin")
panda_path = hf_hub_download(repo_id="openllmplayground/pandagpt_13b_max_len_400", filename="pytorch_model.pt")

a = PandaGPT((vicuna_path,))
a.load_projection(panda_path)


def add_text(history, text):
   history = history + [(text, None)]
   return history, gr.update(value="", interactive=False)

def add_file(history, file):
    history = history + [((file.name,), None)]
    return history


def bot(history):
    text = history[-1][0]
    image_paths = []
    audio_paths = []
    video_paths = []
    for i in history[:-1]:
        if i[1] is None:
            if i[0][:4] in [".png", "jpeg"]:
                image_paths += list(i[0])
            if i[0][:3] in ["mp3", "wav"]:
                audio_paths += list(i[0])
            if i[0][:3] in ["mp4", "avi", "mkv"]:
                video_paths += list(i[0])
        else:
            image_paths = []
            audio_paths = []
            video_paths = []
    if len(image_paths) == 0 and len(audio_paths) == 0 and len(video_paths) == 0:
        response = a.eval_with_image(None, text)
    else:
        response = a.eval_with_image({"image_paths": image_paths,"audio_paths": audio_paths, "video_paths": video_paths}, text)
    history[-1][1] = ""
    for i in a.generate():
        history[-1][1] += i
        yield history
    if history[-1][1].endswith("###"):
        history[-1][1] = history[-1][1][:-3]
    yield history


with gr.Blocks() as demo:
    chatbot = gr.Chatbot([], elem_id="chatbot").style(height=750)

    with gr.Row():
        with gr.Column(scale=0.65):
            txt = gr.Textbox(
                show_label=False,
                placeholder="Enter text and press enter, or upload an image",
            ).style(container=False)
        with gr.Column(scale=0.15, min_width=0):
            btn = gr.UploadButton("📁", file_types=["image", "video", "audio"])
        with gr.Column(scale=0.15, min_width=0):
            btn2 = gr.Button("reset")

    txt_msg = txt.submit(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(
        bot, chatbot, chatbot
    )
    txt_msg.then(lambda: gr.update(interactive=True), None, [txt], queue=False)
    file_msg = btn.upload(add_file, [chatbot, btn], [chatbot], queue=False)
    btn2.click(a.reset, None, chatbot, queue=False)

demo.queue()
demo.launch()


# a.chat_with