from transformers import MllamaForConditionalGeneration, AutoProcessor, TextIteratorStreamer
from PIL import Image
import requests
import torch
from threading import Thread
import gradio as gr
from gradio import FileData
import time
import spaces
ckpt = "meta-llama/Llama-3.2-11B-Vision-Instruct"
model = MllamaForConditionalGeneration.from_pretrained(ckpt,
    torch_dtype=torch.bfloat16).to("cuda")
processor = AutoProcessor.from_pretrained(ckpt)


@spaces.GPU
def bot_streaming(message, history, max_new_tokens=250):
    
    txt = message["text"]
    ext_buffer = f"{txt}"
    
    messages= [] 
    images = []
    

    for i, msg in enumerate(history): 
        if isinstance(msg[0], tuple):
            messages.append({"role": "user", "content": [{"type": "text", "text": history[i+1][0]}, {"type": "image"}]})
            messages.append({"role": "assistant", "content": [{"type": "text", "text": history[i+1][1]}]})
            images.append(Image.open(msg[0][0]).convert("RGB"))
        elif isinstance(history[i-1], tuple) and isinstance(msg[0], str):
            # messages are already handled
            pass
        elif isinstance(history[i-1][0], str) and isinstance(msg[0], str): # text only turn
            messages.append({"role": "user", "content": [{"type": "text", "text": msg[0]}]})
            messages.append({"role": "assistant", "content": [{"type": "text", "text": msg[1]}]})

    # add current message
    if len(message["files"]) == 1:
        
        if isinstance(message["files"][0], str): # examples
            image = Image.open(message["files"][0]).convert("RGB")
        else: # regular input
            image = Image.open(message["files"][0]["path"]).convert("RGB")
        images.append(image)
        messages.append({"role": "user", "content": [{"type": "text", "text": txt}, {"type": "image"}]})
    else:
        messages.append({"role": "user", "content": [{"type": "text", "text": txt}]})


    texts = processor.apply_chat_template(messages, add_generation_prompt=True)

    if images == []:
        inputs = processor(text=texts, return_tensors="pt").to("cuda")
    else:
        inputs = processor(text=texts, images=images, return_tensors="pt").to("cuda")
    streamer = TextIteratorStreamer(processor, skip_special_tokens=True, skip_prompt=True)

    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=max_new_tokens)
    generated_text = ""
    
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    buffer = ""
    
    for new_text in streamer:
        buffer += new_text
        generated_text_without_prompt = buffer
        time.sleep(0.01)
        yield buffer


demo = gr.ChatInterface(fn=bot_streaming, title="Multimodal Llama", examples=[
    [{"text": "Which era does this piece belong to? Give details about the era.", "files":["./examples/rococo.jpg"]},
    200],
    [{"text": "Where do the droughts happen according to this diagram?", "files":["./examples/weather_events.png"]},
    250],
    [{"text": "What happens when you take out white cat from this chain?", "files":["./examples/ai2d_test.jpg"]},
    250],
    [{"text": "How long does it take from invoice date to due date? Be short and concise.", "files":["./examples/invoice.png"]},
    250],
    [{"text": "Where to find this monument? Can you give me other recommendations around the area?", "files":["./examples/wat_arun.jpg"]}, 
    250],
    ],
      textbox=gr.MultimodalTextbox(), 
      additional_inputs = [gr.Slider(
              minimum=10,
              maximum=500,
              value=250,
              step=10,
              label="Maximum number of new tokens to generate",
          )
        ],
      cache_examples=False,
      description="Try Multimodal Llama by Meta with transformers in this demo. Upload an image, and start chatting about it, or simply try one of the examples below. To learn more about Llama Vision, visit [our blog post](https://huggingface.co/blog/llama32). ",
      stop_btn="Stop Generation", 
      fill_height=True,
    multimodal=True)
    
demo.launch(server_name="10.172.0.233", server_port=7860, share=True)

Traceback (most recent call last):
  File "/maahr/home/ameen/anaconda3/envs/llama/lib/python3.10/site-packages/gradio/queueing.py", line 624, in process_events
    response = await route_utils.call_process_api(
  File "/maahr/home/ameen/anaconda3/envs/llama/lib/python3.10/site-packages/gradio/route_utils.py", line 323, in call_process_api
    output = await app.get_blocks().process_api(
  File "/maahr/home/ameen/anaconda3/envs/llama/lib/python3.10/site-packages/gradio/blocks.py", line 2015, in process_api
    result = await self.call_function(
  File "/maahr/home/ameen/anaconda3/envs/llama/lib/python3.10/site-packages/gradio/blocks.py", line 1574, in call_function
    prediction = await utils.async_iteration(iterator)
  File "/maahr/home/ameen/anaconda3/envs/llama/lib/python3.10/site-packages/gradio/utils.py", line 710, in async_iteration
    return await anext(iterator)
  File "/maahr/home/ameen/anaconda3/envs/llama/lib/python3.10/site-packages/gradio/utils.py", line 815, in asyncgen_wrapper
    response = await iterator.__anext__()
  File "/maahr/home/ameen/anaconda3/envs/llama/lib/python3.10/site-packages/gradio/chat_interface.py", line 678, in _stream_fn
    first_response = await async_iteration(generator)
  File "/maahr/home/ameen/anaconda3/envs/llama/lib/python3.10/site-packages/gradio/utils.py", line 710, in async_iteration
    return await anext(iterator)
  File "/maahr/home/ameen/anaconda3/envs/llama/lib/python3.10/site-packages/gradio/utils.py", line 704, in __anext__
    return await anyio.to_thread.run_sync(
  File "/maahr/home/ameen/anaconda3/envs/llama/lib/python3.10/site-packages/anyio/to_thread.py", line 56, in run_sync
    return await get_async_backend().run_sync_in_worker_thread(
  File "/maahr/home/ameen/anaconda3/envs/llama/lib/python3.10/site-packages/anyio/_backends/_asyncio.py", line 2144, in run_sync_in_worker_thread
    return await future
  File "/maahr/home/ameen/anaconda3/envs/llama/lib/python3.10/site-packages/anyio/_backends/_asyncio.py", line 851, in run
    result = context.run(func, *args)
  File "/maahr/home/ameen/anaconda3/envs/llama/lib/python3.10/site-packages/gradio/utils.py", line 687, in run_sync_iterator_async
    return next(iterator)
  File "/maahr/home/ameen/LLM/llama-3.2-vision-11B/app.py", line 42, in bot_streaming
    image = Image.open(message["files"][0]).convert("RGB")
  File "/maahr/home/ameen/anaconda3/envs/llama/lib/python3.10/site-packages/PIL/Image.py", line 3339, in open
    raise UnidentifiedImageError(msg)
PIL.UnidentifiedImageError: cannot identify image file '/tmp/gradio/6cc60d3cbb27ea1fd39b0eab1f4d53124661902c0f942e1461f1c32fd96f8dae/update_Energy_dataset.csv'