Update app.py

app.py

@@ -1,10 +1,7 @@
 import os
 import torch
 import gradio as gr
-import tempfile
-import secrets
-from pathlib import Path
-from transformers import AutoModelForCausalLM, AutoTokenizer, BlipForConditionalGeneration, AutoProcessor, Qwen2VLForConditionalGeneration
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor, Qwen2VLForConditionalGeneration
 from qwen_vl_utils import process_vision_info
 from PIL import Image
 
@@ -12,9 +9,6 @@ from PIL import Image
 vl_model = Qwen2VLForConditionalGeneration.from_pretrained(
     "Qwen/Qwen2-VL-2B-Instruct", torch_dtype="auto", device_map="auto"
 )
-Qwen2VLForConditionalGeneration.from_pretrained(
-    "Qwen/Qwen2-VL-2B-Instruct", torch_dtype="auto", device_map="auto"
-)
 vl_processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct")
 
 # Load Text Model
@@ -25,27 +19,19 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 math_messages = []
 
-def process_image(image, shouldConvert=False):
+def process_image(image, should_convert=False):
     global math_messages
     math_messages = []  # Reset when uploading an image
 
-    if shouldConvert:
+    if should_convert:
         new_img = Image.new('RGB', size=(image.width, image.height), color=(255, 255, 255))
         new_img.paste(image, (0, 0), mask=image)
         image = new_img
 
-    # Convert the image to tensor
-    inputs = vl_processor(images=image, return_tensors="pt")
+    inputs = vl_processor(images=image, return_tensors="pt").to(device)
     generated_ids = vl_model.generate(**inputs)
-    generated_ids_trimmed = [
-        out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
-    ]
-    output = processor.batch_decode(
-        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
-    )
-    description = vl_processor.batch_decode(output, skip_special_tokens=True)[0]
-
-    return f"Math-related content detected: {description}"
+    output = vl_processor.batch_decode(generated_ids, skip_special_tokens=True)
+    return f"Math-related content detected: {output[0]}"
 
 def get_math_response(image_description, user_question):
     global math_messages
@@ -55,21 +41,23 @@ def get_math_response(image_description, user_question):
     content = f'Image description: {image_description}\n\n' if image_description else ''
     query = f"{content}User question: {user_question}"
     math_messages.append({'role': 'user', 'content': query})
+    
     model_inputs = tokenizer(query, return_tensors="pt").to(device)
     output = model.generate(**model_inputs, max_new_tokens=512)
     answer = tokenizer.decode(output[0], skip_special_tokens=True)
+    
     yield answer.replace("\\", "\\\\")
     math_messages.append({'role': 'assistant', 'content': answer})
 
 def math_chat_bot(image, sketchpad, question, state):
     current_tab_index = state["tab_index"]
     image_description = None
-    if current_tab_index == 0:
-        if image is not None:
-            image_description = process_image(image)
-    elif current_tab_index == 1:
-        if sketchpad and sketchpad["composite"]:
-            image_description = process_image(sketchpad["composite"], True)
+    
+    if current_tab_index == 0 and image is not None:
+        image_description = process_image(image)
+    elif current_tab_index == 1 and sketchpad and sketchpad["composite"]:
+        image_description = process_image(sketchpad["composite"], True)
+    
     yield from get_math_response(image_description, question)
 
 css = """
@@ -81,67 +69,41 @@ css = """
 def tabs_select(e: gr.SelectData, _state):
     _state["tab_index"] = e.index
 
-
-# 创建Gradio接口
+# Create Gradio UI
 with gr.Blocks(css=css) as demo:
-    gr.HTML("""\
-<p align="center"><img src="https://modelscope.oss-cn-beijing.aliyuncs.com/resource/qwen.png" style="height: 60px"/><p>"""
-            """<center><font size=8>📖 Qwen2-Math Demo</center>"""
-            """\
-<center><font size=3>This WebUI is based on Qwen2-VL for OCR and Qwen2-Math for mathematical reasoning. You can input either images or texts of mathematical or arithmetic problems.</center>"""
-            )
+    gr.HTML("""
+    <p align="center"><img src="https://modelscope.oss-cn-beijing.aliyuncs.com/resource/qwen.png" style="height: 60px"/></p>
+    <center><font size=8>📖 Qwen2-Math Demo</font></center>
+    <center><font size=3>This WebUI is based on Qwen2-VL for OCR and Qwen2-Math for mathematical reasoning. You can input either images or texts of mathematical or arithmetic problems.</font></center>
+    """)
+    
     state = gr.State({"tab_index": 0})
+    
     with gr.Row():
         with gr.Column():
             with gr.Tabs() as input_tabs:
                 with gr.Tab("Upload"):
-                    input_image = gr.Image(type="pil", label="Upload"),
+                    input_image = gr.Image(type="pil", label="Upload")
                 with gr.Tab("Sketch"):
-                    input_sketchpad = gr.Sketchpad(type="pil", label="Sketch", layers=False)
+                    input_sketchpad = gr.Sketchpad(label="Sketch", layers=False)
+            
             input_tabs.select(fn=tabs_select, inputs=[state])
-            input_text = gr.Textbox(label="input your question")
+            input_text = gr.Textbox(label="Input your question")
+            
             with gr.Row():
                 with gr.Column():
-                    clear_btn = gr.ClearButton(
-                        [*input_image, input_sketchpad, input_text])
+                    clear_btn = gr.ClearButton([input_image, input_sketchpad, input_text])
                 with gr.Column():
                     submit_btn = gr.Button("Submit", variant="primary")
+        
         with gr.Column():
-            output_md = gr.Markdown(label="answer",
-                                    latex_delimiters=[{
-                                        "left": "\\(",
-                                        "right": "\\)",
-                                        "display": True
-                                    }, {
-                                        "left": "\\begin\{equation\}",
-                                        "right": "\\end\{equation\}",
-                                        "display": True
-                                    }, {
-                                        "left": "\\begin\{align\}",
-                                        "right": "\\end\{align\}",
-                                        "display": True
-                                    }, {
-                                        "left": "\\begin\{alignat\}",
-                                        "right": "\\end\{alignat\}",
-                                        "display": True
-                                    }, {
-                                        "left": "\\begin\{gather\}",
-                                        "right": "\\end\{gather\}",
-                                        "display": True
-                                    }, {
-                                        "left": "\\begin\{CD\}",
-                                        "right": "\\end\{CD\}",
-                                        "display": True
-                                    }, {
-                                        "left": "\\[",
-                                        "right": "\\]",
-                                        "display": True
-                                    }],
-                                    elem_id="qwen-md")
+            output_md = gr.Markdown(label="Answer", elem_id="qwen-md")
+        
         submit_btn.click(
             fn=math_chat_bot,
-            inputs=[*input_image, input_sketchpad, input_text, state],
-            outputs=output_md)
-        
+            inputs=[input_image, input_sketchpad, input_text, state],
+            outputs=output_md
+        )
+
 if __name__ == "__main__":
     demo.launch()