import gradio as gr
import transformers
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation import GenerationConfig
import re
import copy
from pathlib import Path
import secrets
import torch
import gc
import os
import io
from io import BytesIO
import base64
import PIL
from PIL import ImageDraw, UnidentifiedImageError
from PIL import Image as PILImage


image_dir = "saved_images"
os.makedirs(image_dir, exist_ok=True)
base_url = "https://huggingface.co/spaces/Tonic1/Official-Qwen-VL-Chat/blob/main/"
model_name = "Qwen/Qwen-VL-Chat"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", bf16=True, trust_remote_code=True).eval()
model.generation_config = GenerationConfig.from_pretrained(model_name, trust_remote_code=True)

BOX_TAG_PATTERN = r"<box>([\s\S]*?)</box>"
PUNCTUATION = "！？。＂＃＄％＆＇（）＊＋，－／：；＜＝＞＠［＼］＾＿｀｛｜｝～｟｠｢｣､、〃》「」『』&#8203;``【oaicite:0】``&#8203;〔〕〖〗〘〙〚〛〜〝〞〟〰〾〿–—‘’‛“”„‟…‧﹏."

class ChatBot:
    def __init__(self, tokenizer, model):
        self.tokenizer = tokenizer
        self.model = model
        self.history = []

    def chat(self, image_path=None, text_query=None):
        query_elements = []
        if image_path:
            query_elements.append({'image': image_path})
        if text_query:
            query_elements.append({'text': text_query})
        
        query = self.tokenizer.from_list_format(query_elements)
        response, self.history = self.model.chat(self.tokenizer, query=query, history=self.history, max_new_tokens = 1200)
        return response

    def draw_boxes(self, response, image_path):
        boxes = re.findall(r'<box>\((\d+),(\d+)\),\((\d+),(\d+)\)</box>', response)
        if not boxes:
            return None
        try:
            with PILImage.open(image_path) as img:
                draw = ImageDraw.Draw(img)
                for box in boxes:
                    x1, y1, x2, y2 = map(int, box)
                    draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
                file_name = secrets.token_hex(10) + ".png"
                file_path = os.path.join(base_url, image_dir, file_name)
                img.save(file_path, format="PNG")
                return file_path
        except Exception as e:
            print(f"An error occurred while processing the image: {e}")
            return None
            
    def clean_response(self, response):
        return re.sub(r'<ref>(.*?)</ref>(?:<box>.*?</box>)*(?:<quad>.*?</quad>)*', r'\1', response).strip()
    
    def clear_memory(self):
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()

def chat_interface(text_query, file):
    chatbot = ChatBot(tokenizer, model)
    image_path = file.name if file is not None else None
    response = chatbot.chat(image_path=image_path, text_query=text_query)

    if "<box>" in response:
        image_file_path = chatbot.draw_boxes(response, image_path)
        text_response = chatbot.clean_response(response)
        chatbot.clear_memory()
        return [("Qwen-VL_Chat", text_response), ("Qwen-VL_Image", image_file_path)]
    else:
        chatbot.clear_memory()
        return [("Qwen-VL_Chat", response)]
        
with gr.Blocks() as demo:
    gr.Markdown("""
# 🙋🏻‍♂️欢迎来到🌟Tonic 的🦆Qwen-VL-Chat🤩Bot！🚀
# 🙋🏻‍♂️Welcome to Tonic's🦆Qwen-VL-Chat🤩Bot！🚀
该WebUI基于Qwen-VL-Chat，实现聊天机器人功能。 但我必须解决它的很多问题，也许我也能获得一些荣誉。
Qwen-VL-Chat 是一种多模式输入模型。 您可以使用此空间来测试当前模型 [qwen/Qwen-VL-Chat](https://huggingface.co/qwen/Qwen-VL-Chat) 您也可以使用 🧑🏻‍🚀qwen/Qwen-VL -通过克隆这个空间来聊天🚀。 🧬🔬🔍 只需点击这里：[重复空间](https://huggingface.co/spaces/Tonic1/VLChat?duplicate=true)
加入我们：🌟TeamTonic🌟总是在制作很酷的演示！ 在 👻Discord 上加入我们活跃的构建者🛠️社区：[Discord](https://discord.gg/nXx5wbX9) 在 🤗Huggingface 上：[TeamTonic](https://huggingface.co/TeamTonic) 和 [MultiTransformer](https:/ /huggingface.co/MultiTransformer) 在 🌐Github 上：[Polytonic](https://github.com/tonic-ai) 并为 🌟 [PolyGPT](https://github.com/tonic-ai/polygpt-alpha) 做出贡献 ）
**You can upload a picture with query or specifically ask this model to draw boxes and frames (just keep trying until you get your returns)** . Qwen-VL-Chat is a multimodal input model. You can use this Space to test out the current model [qwen/Qwen-VL-Chat](https://huggingface.co/qwen/Qwen-VL-Chat) You can also use  qwen/Qwen-VL-Chat🚀 by cloning this space.   Simply click here: [Duplicate Space](https://huggingface.co/spaces/Tonic1/VLChat?duplicate=true)
Join us:  TeamTonic  is always making cool demos! Join our active builder's community on  Discord: [Discord](https://discord.gg/nXx5wbX9) On  Huggingface: [TeamTonic](https://huggingface.co/TeamTonic) & [MultiTransformer](https://huggingface.co/MultiTransformer) On  Github: [Polytonic](https://github.com/tonic-ai) & contribute to   [PolyGPT](https://github.com/tonic-ai/polygpt-alpha)
""")
    with gr.Row():
        with gr.Column(scale=1):
            chatbot_component = gr.Chatbot(label='🦆Qwen-VL-Chat')
        with gr.Column(scale=1):
            with gr.Row():
                query = gr.Textbox(lines=2, label='Input', placeholder="Type your message here...")
                file_upload = gr.File(label="Upload Image")
                submit_btn = gr.Button("Submit")
    
    submit_btn.click(
        fn=chat_interface,
        inputs=[query, file_upload],
        outputs=[chatbot_component]
    )
    gr.Markdown("""
注意：此演示受 Qwen-VL 原始许可证的约束。我们强烈建议用户不要故意生成或允许他人故意生成有害内容，
包括仇恨言论、暴力、色情、欺骗等。（注：本演示受Qwen-VL许可协议约束，强烈建议用户不要传播或允许他人传播以下内容，包括但不限于仇恨言论、暴力、色情、欺诈相关的有害信息 .)
Note: This demo is governed by the original license of Qwen-VL. We strongly advise users not to knowingly generate or allow others to knowingly generate harmful content,
including hate speech, violence, pornography, deception, etc. (Note: This demo is subject to the license agreement of Qwen-VL. We strongly advise users not to disseminate or allow others to disseminate the following content, including but not limited to hate speech, violence, pornography, and fraud-related harmful information.)
""")
if __name__ == "__main__":
    demo.launch()