Update app.py

app.py

@@ -15,7 +15,7 @@ def run_GOT(image_array, got_mode, ocr_box="", ocr_color=""):
     # image = Image.fromarray(np.uint8(image_array))
     image = image_array
     if got_mode == "plain texts OCR":
-        res = model.chat(tokenizer, image, ocr_type='ocr')
+        res = model.chat(tokenizer, image, ocr_type='ocr', stream_flag = True)
     elif got_mode == "format texts OCR":
         res = model.chat(tokenizer, image, ocr_type='format', render=True, save_render_file='./demo.html')
     elif got_mode == "plain multi-crop OCR":
@@ -31,8 +31,8 @@ def run_GOT(image_array, got_mode, ocr_box="", ocr_color=""):
     if "format" in got_mode:
         with open('./demo.html', 'r') as f:
             demo_html = f.read()
-        return res, demo_html
-    return res, None
+        yield res, demo_html
+    yield res, None
 
 def task_update(task):
     if "fine-grained" in task:
@@ -63,7 +63,7 @@ def fine_grained_update(task):
 
 with gr.Blocks() as demo:
     gr.Markdown("""
-    # "General OCR Theory: Towards OCR-2.0 via a Unified End-to-end Model"
+    # General OCR Theory: Towards OCR-2.0 via a Unified End-to-end Model
     
     "🔥🔥🔥This is the official online demo of GOT-OCR-2.0 model!!!"
     
@@ -106,7 +106,8 @@ with gr.Blocks() as demo:
             submit_button = gr.Button("Submit")
         
         with gr.Column():
-            ocr_result = gr.Textbox(label="GOT output")
+            chatbot = gr.Chatbot(label="GOT output")
+            # ocr_result = gr.Textbox(label="GOT output")
             html_result = gr.HTML(label="rendered html")
     
     gr.Examples(
@@ -120,7 +121,7 @@ with gr.Blocks() as demo:
             ["assets/color.png", "plain fine-grained OCR", "color", "red", ""],
         ],
         inputs=[image_input, task_dropdown, fine_grained_dropdown, color_dropdown, box_input],
-        outputs=[ocr_result, html_result],
+        outputs=[chatbot, html_result],
         label="examples",
     )
 
@@ -138,7 +139,7 @@ with gr.Blocks() as demo:
     submit_button.click(
         run_GOT,
         inputs=[image_input, task_dropdown, box_input, color_dropdown],
-        outputs=[ocr_result, html_result]
+        outputs=[chatbot, html_result]
     )
 
 demo.launch(share=True)