import spaces
import argparse
import torch
import re
import gradio as gr
from threading import Thread
from transformers import TextIteratorStreamer, AutoTokenizer, AutoModelForCausalLM
from PIL import Image

parser = argparse.ArgumentParser()

model_id = "vikhyat/moondream2"
revision = "2024-04-02"
tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision)
moondream = AutoModelForCausalLM.from_pretrained(
    model_id, trust_remote_code=True, revision=revision,
    torch_dtype=torch.float32
)
moondream.eval()

@spaces.GPU(duration=10)
def answer_question(images, prompts):
    image_embeds = [moondream.encode_image(img) for img in images]
    image_embeds = torch.cat(image_embeds, dim=0)
    answers = moondream.batch_answer(
        images=image_embeds,
        prompts=prompts,
        tokenizer=tokenizer
    )
    return [answer for answer in answers]

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🌔 moondream2
        A tiny vision language model. [GitHub](https://github.com/vikhyat/moondream)
        """
    )
    with gr.Row():
        prompts = gr.Textbox(label="Input", placeholder="Type here...", scale=4)
        submit = gr.Button("Submit")
    with gr.Row():
        images = gr.Image(type="pil", label="Upload Images", multiple=True)
        output = gr.Textbox(label="Response", multiple=True)
    submit.click(answer_question, [images, prompts], output)
    prompts.submit(answer_question, [images, prompts], output)

demo.queue().launch()