Update app.py

app.py

@@ -26,112 +26,54 @@ task_history = []
 BOX_TAG_PATTERN = r"<box>([\s\S]*?)</box>"
 PUNCTUATION = "！？。＂＃＄％＆＇（）＊＋，－／：；＜＝＞＠［＼］＾＿｀｛｜｝～｟｠｢｣､、〃》「」『』&#8203;``【oaicite:0】``&#8203;〔〕〖〗〘〙〚〛〜〝〞〟〰〾〿–—‘’‛“”„‟…‧﹏."
 
-import os
-from PIL import Image, UnidentifiedImageError
-import secrets
-
-
-def save_image(image_file) -> str:
-    upload_dir = os.path.abspath("uploaded_images")
-    print(f"Creating upload directory at {upload_dir} if it doesn't exist.")
-    os.makedirs(upload_dir, exist_ok=True)
-    try:
-        image = Image.open(image_file).convert("RGB")
-        file_name = secrets.token_hex(10) + ".png"
-        file_path = os.path.join(upload_dir, file_name)
-        print(f"Generated file path: {file_path}")
-        print("Saving the image.")
-        image.save(file_path, format="PNG")
-        print("Image saved successfully.")
-        return file_path
-    except UnidentifiedImageError:
-        print("Error: The file is not a recognizable image.")
-        return None
-    except Exception as e:
-        print(f"An unexpected error occurred: {e}")
-        return None
+class ChatBot:
+    def __init__(self):
+        self.tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen-VL-Chat", trust_remote_code=True)
+        self.model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-VL-Chat", device_map="cuda", trust_remote_code=True).eval()
+        self.model.generation_config = GenerationConfig.from_pretrained("Qwen/Qwen-VL-Chat", trust_remote_code=True)
+        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})
         
-def clean_response(response: str) -> str:
-    response = re.sub(r'<ref>(.*?)</ref>(?:<box>.*?</box>)*(?:<quad>.*?</quad>)*', r'\1', response).strip()
-    return response
-
-def chat_with_model(image_path=None, text_query=None, history=None):
-    if image_path:
-        try:
-            with Image.open(image_path) as img:
-                print(f"Image {image_path} opened successfully.")
-        except UnidentifiedImageError:
-            print(f"Error: The file at {image_path} is not a recognizable image.")
-            return "Error: Uploaded file is not a recognizable image."
-        except Exception as e:
-            print(f"An error occurred while processing the image: {e}")
-            return "An error occurred while processing your request."
-    else:
-        print("No image path provided, using text-only mode.")
-
-    text_input = text_query if text_query else ""
-    query_elements = [
-        {'image': image_path},
-        {'text': text_input}
-    ]
-    try:
-        query = tokenizer.from_list_format(query_elements)
-        tokenized_inputs = tokenizer(query, return_tensors='pt').to(device)
-        output = model.generate(**tokenized_inputs)
-        response = tokenizer.decode(output[0], skip_special_tokens=True)
-        cleaned_response = clean_response(response)
-        return cleaned_response
-    except Exception as e:
-        print(f"An error occurred: {e}")
-        return "An error occurred while processing your request."
-
-def draw_boxes(image_path, response):
-    with Image.open(image_path) as image:
-        draw = ImageDraw.Draw(image)
-        boxes = re.findall(r'<box>\((\d+),(\d+)\),\((\d+),(\d+)\)</box>', response)
-        for box in boxes:
-            x1, y1, x2, y2 = map(int, box)
-            draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
-        with tempfile.NamedTemporaryFile(delete=False, suffix='.png', dir=os.path.abspath("uploaded_images")) as temp_file:
-            image.save(temp_file, format="PNG")
-            return temp_file.name
-
-def draw_boxes_with_tokenizer(response, history):
-    image = tokenizer.draw_bbox_on_latest_picture(response, history)
-    if image is not None:
-        buffered = io.BytesIO()
-        image.save(buffered, format="PNG")
-        img_str = base64.b64encode(buffered.getvalue()).decode()
-        return "data:image/png;base64," + img_str
-    else:
-        print("No box found or drawing failed.")
-        return None
+        query = self.tokenizer.from_list_format(query_elements)
+        response, self.history = self.model.chat(self.tokenizer, query=query, history=self.history)
+        return response
+
+    def draw_boxes(self, response):
+        image = self.tokenizer.draw_bbox_on_latest_picture(response, self.history)
+        if image is not None:
+            buffered = io.BytesIO()
+            image.save(buffered, format="PNG")
+            img_str = base64.b64encode(buffered.getvalue()).decode()
+            return "data:image/png;base64," + img_str
+        else:
+            return None
 
-def process_input(text=None, file=None, task_history=None):
-    if task_history is None:
-        task_history = []
-    image_path = None
-    if file is not None:
-        image_path = save_image(file)
-        if image_path is None:
-            return [("bot", "Error: Uploaded file is not a recognizable image.")], task_history
+    def clear_memory(self):
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        gc.collect()
 
-    response = chat_with_model(image_path=image_path, text_query=text, history=task_history)
-    task_history.append((text, response))
+def chat_interface(chatbot, text_query, file):
+    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 and image_path:
-        image_with_boxes_base64 = draw_boxes_with_tokenizer(response, task_history)
-        if image_with_boxes_base64:
-            cleaned_response = clean_response(response)
-            return [("Qwen-VL_Image", image_with_boxes_base64), ("Qwen-VL_Chat", cleaned_response)], task_history
-        else:
-            return [("bot", "Unable to draw boxes on the image.")], task_history
+    if "<box>" in response:
+        image_with_boxes = chatbot.draw_boxes(response)
+        chatbot.clear_memory()
+        return [("Qwen-VL_Chat", response), ("Qwen-VL_Image", image_with_boxes)]
     else:
-        cleaned_response = clean_response(response)
-        return [("Qwen-VL_Chat", cleaned_response)], task_history
+        chatbot.clear_memory()
+        return [("Qwen-VL_Chat", response)]
 
 with gr.Blocks() as demo:
-    gr.Markdown("""
+    chatbot = ChatBot()
+    with gr.Markdown("""
 # 🙋🏻‍♂️欢迎来到🌟Tonic 的🦆Qwen-VL-Chat🤩Bot！🚀
 # 🙋🏻‍♂️Welcome to Tonic's🦆Qwen-VL-Chat🤩Bot！🚀
 该WebUI基于Qwen-VL-Chat，实现聊天机器人功能。 但我必须解决它的很多问题，也许我也能获得一些荣誉。
@@ -142,28 +84,23 @@ Join us:  TeamTonic  is always making cool demos! Join our active builder's comm
 """)
     with gr.Row():
         with gr.Column(scale=1):
-            chatbot = gr.Chatbot(label='Qwen-VL-Chat')
+            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.Image(type="filepath", label="Upload Image")
+                file_upload = gr.File(label="Upload Image")
                 submit_btn = gr.Button("Submit")
     
-    task_history = gr.State([])
-
     submit_btn.click(
-        fn=process_input,
-        inputs=[query, file_upload, task_history],
-        outputs=[chatbot, task_history]
+        fn=chat_interface,
+        inputs=[chatbot, query, file_upload],
+        outputs=[chatbot_component]
     )
-    
-    gr.Markdown("""
+    with 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.)
 """)
-    demo.queue().launch()
-
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()