import gradio as gr
import torch

from threading import Thread
from PIL import Image
from transformers import TextIteratorStreamer
from transformers import LlavaNextForConditionalGeneration, LlavaNextProcessor
from PIL import Image

import spaces

PLACEHOLDER = """
<div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
   <img src="https://raw.githubusercontent.com/huggingface/blog/09dbdfd196a3112ecbb533fc0b6c700571cbc753/assets/179_falcon2-11b/thumbnail.jpg" style="width: 80%; max-width: 550px; height: auto; opacity: 0.55;  "> 
   <h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">Falcon2-11B-VLM</h1>
   <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.65;">Falcon2-11B-VLM is an 11B parameters causal decoder-only model built by TII</p>
</div>
"""
model_id = "tiiuae/falcon-11B-vlm"

processor = LlavaNextProcessor.from_pretrained("tiiuae/falcon-11B-vlm", tokenizer_class='PreTrainedTokenizerFast')
model = LlavaNextForConditionalGeneration.from_pretrained("tiiuae/falcon-11B-vlm", 
                                                          torch_dtype=torch.bfloat16,
                                                          #torch_dtype=torch.float16,
                                                          low_cpu_mem_usage=True,)

model.to("cuda:0")

@spaces.GPU
def bot_streaming(message, history):
    print(f'message is - {message}')
    print(f'history is - {history}')
    if message["files"]:
        # message["files"][-1] is a Dict or just a string
        if type(message["files"][-1]) == dict:
            image = message["files"][-1]["path"]
        else:
            image = message["files"][-1]
    else:
        # if there's no image uploaded for this turn, look for images in the past turns
        # kept inside tuples, take the last one
        for hist in history:
            if type(hist[0]) == tuple:
                image = hist[0][0]
    try:
        if image is None:
            # Handle the case where image is None
            raise gr.Error("You need to upload an image for FalconVLM to work. Close the error and try again with an Image.")
    except NameError:
        # Handle the case where 'image' is not defined at all
        raise gr.Error("You need to upload an image for FalconVLM to work. Close the error and try again with an Image.")

    prompt = f"""User:<image>\n{message['text']} Falcon:"""
    image = Image.open(image)
    inputs = processor(prompt, image, return_tensors='pt').to(0, torch.float16)

    streamer = TextIteratorStreamer(processor, **{"skip_special_tokens": True, "skip_prompt": True})
    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024, do_sample=False)

    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    buffer = ""
    for new_text in streamer:
        buffer += new_text
        yield buffer


chatbot=gr.Chatbot(placeholder=PLACEHOLDER,scale=1)
chat_input = gr.MultimodalTextbox(interactive=True, file_types=["image"], placeholder="Enter message or upload file...", show_label=False)
with gr.Blocks(fill_height=True, ) as demo:
    gr.ChatInterface(
    fn=bot_streaming,
    title="FalconVLM",
    examples=[{"text": "What is on the flower?", "files": ["./bee.jpg"]},
              {"text": "How to make this pastry?", "files": ["./baklava.png"]}],
    description="Try [tiiuae/falcon-11B-VLM](https://huggingface.co/tiiuae/falcon-11B-vlm). Upload an image and start chatting about it, or simply try one of the examples below. If you don't upload an image, you will receive an error. This is not the official demo.",
    stop_btn="Stop Generation",
    multimodal=True,
    textbox=chat_input,
    chatbot=chatbot,
    cache_examples=False,
    )

demo.queue()
demo.launch()