Update app.py

app.py

@@ -5,64 +5,125 @@ from PIL import Image
 import re 
 import requests
 from io import BytesIO
+import copy
+import secrets
+from pathlib import Path
 
 tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen-VL-Chat-Int4", trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-VL-Chat-Int4", device_map="auto", trust_remote_code=True).eval()
 
-def generate_predictions(image_input, text_input, with_grounding):
-    user_image_path = "/tmp/user_input_test_image.jpg"
-    original_image = Image.fromarray((255 - (image_input * 255).astype('uint8')))
-    original_image.save(user_image_path)
-    
-    if with_grounding == "Yes":
-        text_input += " with grounding"
-    
-    query = tokenizer.from_list_format([
-        {'image': user_image_path},
-        {'text': text_input},
-    ])
-    inputs = tokenizer(query, return_tensors='pt')
-    inputs = inputs.to(model.device)
-    
-    pred = model.generate(**inputs)
-    full_response = tokenizer.decode(pred.cpu()[0], skip_special_tokens=False)
-    
-    frontend_response = re.sub(r'Picture \d+:|<.*?>|\/tmp\/.*\.jpg', '', full_response).replace(text_input, '').strip()
-    print("Generated Caption:", frontend_response)  # Debugging line
-    
-    image_with_boxes = tokenizer.draw_bbox_on_latest_picture(full_response)
-    
-    # Check if the response contains bounding box coordinates
-    if not re.search(r'\(\d+,\d+\),\(\d+,\d+\)', frontend_response):
-        image_with_boxes = original_image
-    
-    if image_with_boxes:
-        temp_path = "/tmp/image_with_boxes.jpg"
-        image_with_boxes.save(temp_path)
-        image_with_boxes = Image.open(temp_path)
-    
-    return image_with_boxes, frontend_response
+BOX_TAG_PATTERN = r"<box>([\s\S]*?)</box>"
+PUNCTUATION = "！？。＂＃＄％＆＇（）＊＋，－／：；＜＝＞＠［＼］＾＿｀｛｜｝～｟｠｢｣､、〃》「」『』&#8203;``oaicite:{"number":1,"invalid_reason":"Malformed citation 【】"}``&#8203;〔〕〖〗〘〙〚〛〜〝〞〟〰〾〿–—‘’‛“”„‟…‧﹏."
+
+def _parse_text(text):
+    lines = text.split("\n")
+    lines = [line for line in lines if line != ""]
+    count = 0
+    for i, line in enumerate(lines):
+        if "```" in line:
+            count += 1
+            items = line.split("`")
+            if count % 2 == 1:
+                lines[i] = f'<pre><code class="language-{items[-1]}">'
+            else:
+                lines[i] = f"<br></code></pre>"
+        else:
+            if i > 0:
+                if count % 2 == 1:
+                    line = line.replace("`", r"\`")
+                    line = line.replace("<", "&lt;")
+                    line = line.replace(">", "&gt;")
+                    line = line.replace(" ", "&nbsp;")
+                    line = line.replace("*", "&ast;")
+                    line = line.replace("_", "&lowbar;")
+                    line = line.replace("-", "&#45;")
+                    line = line.replace(".", "&#46;")
+                    line = line.replace("!", "&#33;")
+                    line = line.replace("(", "&#40;")
+                    line = line.replace(")", "&#41;")
+                    line = line.replace("$", "&#36;")
+                lines[i] = "<br>" + line
+    text = "".join(lines)
+    return text
 
-iface = gr.Interface(
-    fn=generate_predictions, 
-    inputs=[
-        gr.inputs.Image(label="Image Input"), 
-        gr.inputs.Textbox(default="Generate a caption for that image:", label="Prompt"),
-        gr.inputs.Radio(["No", "Yes"], label="With Grounding", default="No")
-    ], 
-    outputs=[
-        gr.outputs.Image(type='pil', label="Image"),
-        gr.outputs.Textbox(label="Generated")
-    ],
-    title="Qwen-VL Demonstration",
-    description = """
-## Qwen-VL: A Multimodal Large Vision Language Model by Alibaba Cloud
-**Space by [@Artificialguybr](https://twitter.com/artificialguybr)**
+def predict(_chatbot, task_history):
+    chat_query = _chatbot[-1][0]
+    query = task_history[-1][0]
+    history_cp = copy.deepcopy(task_history)
+    full_response = ""
+
+    history_filter = []
+    pic_idx = 1
+    pre = ""
+    for i, (q, a) in enumerate(history_cp):
+        if isinstance(q, (tuple, list)):
+            q = f'Picture {pic_idx}: <img>{q[0]}</img>'
+            pre += q + '\n'
+            pic_idx += 1
+        else:
+            pre += q
+            history_filter.append((pre, a))
+            pre = ""
+    history, message = history_filter[:-1], history_filter[-1][0]
+    response, history = model.chat(tokenizer, message, history=history)
+    image = tokenizer.draw_bbox_on_latest_picture(response, history)
+    if image is not None:
+        temp_dir = secrets.token_hex(20)
+        temp_dir = Path("/tmp") / temp_dir
+        temp_dir.mkdir(exist_ok=True, parents=True)
+        name = f"tmp{secrets.token_hex(5)}.jpg"
+        filename = temp_dir / name
+        image.save(str(filename))
+        _chatbot[-1] = (_parse_text(chat_query), (str(filename),))
+        chat_response = response.replace("<ref>", "")
+        chat_response = chat_response.replace(r"</ref>", "")
+        chat_response = re.sub(BOX_TAG_PATTERN, "", chat_response)
+        if chat_response != "":
+            _chatbot.append((None, chat_response))
+    else:
+        _chatbot[-1] = (_parse_text(chat_query), response)
+    full_response = _parse_text(response)
+    task_history[-1] = (query, full_response)
+    return _chatbot
     
-### Key Features:
-- **Strong Performance**: Surpasses existing LVLMs on multiple English benchmarks including Zero-shot Captioning and VQA.
-- **Multi-lingual Support**: Supports English, Chinese, and multi-lingual conversation.
-- **High Resolution**: Utilizes 448*448 resolution for fine-grained recognition and understanding.
-""",
-)
-iface.launch()
+def add_text(history, task_history, text):
+    task_text = text
+    if len(text) >= 2 and text[-1] in PUNCTUATION and text[-2] not in PUNCTUATION:
+        task_text = text[:-1]
+    history = history + [(_parse_text(text), None)]
+    task_history = task_history + [(task_text, None)]
+    return history, task_history, ""
+
+def add_file(history, task_history, file):
+    history = history + [((file.name,), None)]
+    task_history = task_history + [((file.name,), None)]
+    return history, task_history
+
+def reset_user_input():
+    return gr.update(value="")
+
+def reset_state(task_history):
+    task_history.clear()
+    return []
+
+with gr.Blocks() as demo:
+    gr.Markdown("<center><font size=8>Qwen-VL-Chat Bot</center>")
+    gr.Markdown("<center><font size=3>Qwen-VL: A Multimodal Large Vision Language Model by Alibaba Cloud **Space by [@Artificialguybr](https://twitter.com/artificialguybr)</center>")
+    gr.Markdown("### Key Features:\n- **Strong Performance**: Surpasses existing LVLMs on multiple English benchmarks including Zero-shot Captioning and VQA.\n- **Multi-lingual Support**: Supports English, Chinese, and multi-lingual conversation.\n- **High Resolution**: Utilizes 448*448 resolution for fine-grained recognition and understanding.")
+    chatbot = gr.Chatbot(label='Qwen-VL-Chat', elem_classes="control-height", height=750)
+    query = gr.Textbox(lines=2, label='Input')
+    task_history = gr.State([])
+
+    with gr.Row():
+        empty_bin = gr.Button("🧹 Clear History")
+        submit_btn = gr.Button("🚀 Submit")
+        regen_btn = gr.Button("🤔️ Regenerate")
+        addfile_btn = gr.UploadButton("📁 Upload", file_types=["image"])
+
+    submit_btn.click(add_text, [chatbot, task_history, query], [chatbot, task_history]).then(
+        predict, [chatbot, task_history], [chatbot], show_progress=True
+    )
+    submit_btn.click(reset_user_input, [], [query])
+    empty_bin.click(reset_state, [task_history], [chatbot], show_progress=True)
+
+demo.launch()