Update app.py

app.py

@@ -16,7 +16,7 @@ BOX_TAG_PATTERN = r"<box>([\s\S]*?)</box>"
 PUNCTUATION = "！？。＂＃＄％＆＇（）＊＋，－／：；＜＝＞＠［＼］＾＿｀｛｜｝～｟｠｢｣､、〃》「」『』【】〔〕〖〗〘〙〚〛〜〝〞〟〰〾〿–—‘’‛“”„‟…‧﹏."
 uploaded_file_dir = os.environ.get("GRADIO_TEMP_DIR") or str(Path(tempfile.gettempdir()) / "gradio")
 
-def _get_args():
+def _get_args() -> ArgumentParser:
     parser = ArgumentParser()
     parser.add_argument("-c", "--checkpoint-path", type=str, default=DEFAULT_CKPT_PATH,
                         help="Checkpoint name or path, default to %(default)r")
@@ -35,7 +35,7 @@ def _get_args():
     args = parser.parse_args()
     return args
 
-def handle_image_submission(_chatbot, task_history, file):
+def handle_image_submission(_chatbot, task_history, file) -> tuple:
     print("handle_image_submission called")
     if file is None:
         print("No file uploaded")
@@ -49,7 +49,7 @@ def handle_image_submission(_chatbot, task_history, file):
     return predict(_chatbot, task_history)
 
     
-def _load_model_tokenizer(args):
+def _load_model_tokenizer(args) -> tuple:
     model_id = args.checkpoint_path
     model_dir = snapshot_download(model_id, revision=args.revision)
     tokenizer = AutoTokenizer.from_pretrained(
@@ -75,7 +75,7 @@ def _load_model_tokenizer(args):
     return model, tokenizer
 
 
-def _parse_text(text):
+def _parse_text(text: str) -> str:
     lines = text.split("\n")
     lines = [line for line in lines if line != ""]
     count = 0
@@ -106,7 +106,7 @@ def _parse_text(text):
     text = "".join(lines)
     return text
 
-def save_image(image_file, upload_dir):
+def save_image(image_file, upload_dir: str) -> str:
     print("save_image called with:", image_file)
     Path(upload_dir).mkdir(parents=True, exist_ok=True)
     filename = secrets.token_hex(10) + Path(image_file.name).suffix
@@ -125,104 +125,106 @@ def add_file(history, task_history, file):
     task_history = task_history + [((file_path,), None)]
     return history, task_history
 
-def _launch_demo(args, model, tokenizer):
-    uploaded_file_dir = os.environ.get("GRADIO_TEMP_DIR") or str(
-        Path(tempfile.gettempdir()) / "gradio"
-    )
-    def predict(_chatbot, task_history):
-        print("predict called")
-        if not _chatbot:
-            return _chatbot 
-        chat_query = _chatbot[-1][0]
-        print("Chat query:", chat_query)
-
-        if isinstance(chat_query, tuple):
-            query = [{'image': chat_query[0]}]
-        else:
-            query = [{'text': _parse_text(chat_query)}]
-
-        print("Query for model:", query)
-        inputs = tokenizer.from_list_format(query)
-        tokenized_inputs = tokenizer(inputs, return_tensors='pt')
-        tokenized_inputs = tokenized_inputs.to(model.device)
-
-        pred = model.generate(**tokenized_inputs)
-        response = tokenizer.decode(pred.cpu()[0], skip_special_tokens=False)
-        print("Model response:", response)
-        if 'image' in query[0]:
-            image = tokenizer.draw_bbox_on_latest_picture(response)
-            if image is not None:
-                image_path = save_image(image, uploaded_file_dir)
-                _chatbot[-1] = (chat_query, (image_path,))
-            else:
-                _chatbot[-1] = (chat_query, "No image to display.")
+
+def predict(_chatbot, task_history) -> list:
+    print("predict called")
+    if not _chatbot:
+        return _chatbot 
+    chat_query = _chatbot[-1][0]
+    print("Chat query:", chat_query)
+
+    if isinstance(chat_query, tuple):
+        query = [{'image': chat_query[0]}]
+    else:
+        query = [{'text': _parse_text(chat_query)}]
+
+    print("Query for model:", query)
+    inputs = tokenizer.from_list_format(query)
+    tokenized_inputs = tokenizer(inputs, return_tensors='pt')
+    tokenized_inputs = tokenized_inputs.to(model.device)
+
+    pred = model.generate(**tokenized_inputs)
+    response = tokenizer.decode(pred.cpu()[0], skip_special_tokens=False)
+    print("Model response:", response)
+    if 'image' in query[0]:
+        image = tokenizer.draw_bbox_on_latest_picture(response)
+        if image is not None:
+            image_path = save_image(image, uploaded_file_dir)
+           _chatbot[-1] = (chat_query, (image_path,))
         else:
-            _chatbot[-1] = (chat_query, response)
+            _chatbot[-1] = (chat_query, "No image to display.")
+    else:
+        _chatbot[-1] = (chat_query, response)
+    return _chatbot
+
+def save_uploaded_image(image_file, upload_dir):
+    if image is None:
+        return None
+    temp_dir = secrets.token_hex(20)
+    temp_dir = Path(uploaded_file_dir) / 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))
+    return str(filename)
+
+def regenerate(_chatbot, task_history) -> list:
+    if not task_history:
         return _chatbot
+    item = task_history[-1]
+    if item[1] is None:
+        return _chatbot
+    task_history[-1] = (item[0], None)
+    chatbot_item = _chatbot.pop(-1)
+    if chatbot_item[0] is None:
+        _chatbot[-1] = (_chatbot[-1][0], None)
+    else:
+        _chatbot.append((chatbot_item[0], None))
+    return predict(_chatbot, task_history)
 
-    def save_uploaded_image(image_file, upload_dir):
-        if image is None:
-            return None
-        temp_dir = secrets.token_hex(20)
-        temp_dir = Path(uploaded_file_dir) / 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))
-        return str(filename)
-
-    def regenerate(_chatbot, task_history):
-        if not task_history:
-            return _chatbot
-        item = task_history[-1]
-        if item[1] is None:
-            return _chatbot
-        task_history[-1] = (item[0], None)
-        chatbot_item = _chatbot.pop(-1)
-        if chatbot_item[0] is None:
-            _chatbot[-1] = (_chatbot[-1][0], None)
-        else:
-            _chatbot.append((chatbot_item[0], None))
-        return predict(_chatbot, task_history)
-
-    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):
-        if file is None:
-            return history, task_history  # Return if no file is uploaded
-        file_path = file.name
-        history = history + [((file.name,), None)]
-        task_history = task_history + [((file.name,), None)]
-        return history, task_history
+def add_text(history, task_history, text) -> tuple:
+    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 reset_user_input():
-        return gr.update(value="")
+def add_file(history, task_history, file):
+    if file is None:
+        return history, task_history  # Return if no file is uploaded
+    file_path = file.name
+    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 process_response(response):
-        response = response.replace("<ref>", "").replace(r"</ref>", "")
-        response = re.sub(BOX_TAG_PATTERN, "", response)
-        return response
-    def process_history_for_model(task_history):
-        processed_history = []
-        for query, response in task_history:
-            if isinstance(query, tuple): 
-                query = {'image': query[0]}
-            else:
-                query = {'text': query}
-            response = response or ""
-            processed_history.append((query, response))
-        return processed_history
+def process_response(response: str) -> str:
+    response = response.replace("<ref>", "").replace(r"</ref>", "")
+    response = re.sub(BOX_TAG_PATTERN, "", response)
+    return response
+def process_history_for_model(task_history) -> list:
+    processed_history = []
+    for query, response in task_history:
+        if isinstance(query, tuple): 
+            query = {'image': query[0]}
+        else:
+            query = {'text': query}
+        response = response or ""
+        processed_history.append((query, response))
+    return processed_history
+
+def reset_state(task_history) -> list:
+    task_history.clear()
+    return []
 
-    def reset_state(task_history):
-        task_history.clear()
-        return []
 
+def _launch_demo(args, model, tokenizer):
+    uploaded_file_dir = os.environ.get("GRADIO_TEMP_DIR") or str(
+        Path(tempfile.gettempdir()) / "gradio"
+    )
 
     with gr.Blocks() as demo:
         gr.Markdown("""# Welcome to Tonic's Qwen-VL-Chat Bot""")