Update app.py

app.py

@@ -1,111 +1,20 @@
-
-from modelscope import AutoModelForCausalLM, AutoTokenizer, GenerationConfig, snapshot_download
-from argparse import ArgumentParser
-from pathlib import Path
-import shutil
-import copy
 import gradio as gr
-import os
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers.generation import GenerationConfig
 import re
+from pathlib import Path
 import secrets
-import tempfile
-
-#GlobalVariables
-os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
-DEFAULT_CKPT_PATH = 'qwen/Qwen-VL-Chat'
-REVISION = 'v1.0.4'
-BOX_TAG_PATTERN = r"<box>([\s\S]*?)</box>"
-PUNCTUATION = "！？。＂＃＄％＆＇（）＊＋，－／：；＜＝＞＠［＼］＾＿｀｛｜｝～｟｠｢｣､、〃》「」『』【】〔〕〖〗〘〙〚〛〜〝〞〟〰〾〿–—‘’‛“”„‟…‧﹏."
-uploaded_file_dir = os.environ.get("GRADIO_TEMP_DIR") or str(Path(tempfile.gettempdir()) / "gradio")
-tokenizer = None
-model = None
-
-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")
-    parser.add_argument("--revision", type=str, default=REVISION)
-    parser.add_argument("--cpu-only", action="store_true", help="Run demo with CPU only")
-
-    parser.add_argument("--share", action="store_true", default=False,
-                        help="Create a publicly shareable link for the interface.")
-    parser.add_argument("--inbrowser", action="store_true", default=False,
-                        help="Automatically launch the interface in a new tab on the default browser.")
-    parser.add_argument("--server-port", type=int, default=8000,
-                        help="Demo server port.")
-    parser.add_argument("--server-name", type=str, default="127.0.0.1",
-                        help="Demo server name.")
+import torch
 
-    args = parser.parse_args()
-    return args
+# Initialize the model and tokenizer
+model_name = "qwen/Qwen-VL-Chat"
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True).eval()
+model.generation_config = GenerationConfig.from_pretrained(model_name, trust_remote_code=True)
 
-def handle_image_submission(_chatbot, task_history, file) -> tuple:
-    if file is None:
-        return _chatbot, task_history
-    file_path = save_image(file, uploaded_file_dir)
-    history_item = ((file_path,), None)
-    _chatbot.append(history_item)
-    task_history.append(history_item)
-    return predict(_chatbot, task_history)
-
-    
-def _load_model_tokenizer(args) -> tuple:
-    global tokenizer, model
-    model_id = args.checkpoint_path
-    model_dir = snapshot_download(model_id, revision=args.revision)
-    tokenizer = AutoTokenizer.from_pretrained(
-        model_dir, trust_remote_code=True, resume_download=True,
-    )
-
-    if args.cpu_only:
-        device_map = "cpu"
-    else:
-        device_map = "auto"
-
-    model = AutoModelForCausalLM.from_pretrained(
-        model_dir,
-        device_map=device_map,
-        trust_remote_code=True,
-        bf16=True,
-        resume_download=True,
-    ).eval()
-    model.generation_config = GenerationConfig.from_pretrained(
-        model_dir, trust_remote_code=True, resume_download=True,
-    )
-
-    return model, tokenizer
-
-
-def _parse_text(text: str) -> str:
-    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
+# Set device for model
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
 
 def save_image(image_file, upload_dir: str) -> str:
     Path(upload_dir).mkdir(parents=True, exist_ok=True)
@@ -115,131 +24,33 @@ def save_image(image_file, upload_dir: str) -> str:
         f_output.write(f_input.read())
     return str(file_path)
 
-
-def add_file(history, task_history, file):
-    if file is None:
-        return history, task_history
-    file_path = save_image(file)
-    history = history + [((file_path,), None)]
-    task_history = task_history + [((file_path,), None)]
-    return history, task_history
-
-def predict(_chatbot, task_history) -> tuple:
-    if not _chatbot:
-        return _chatbot, task_history
-    chat_query, chat_response = _chatbot[-1]
-    print("predict called")
-    if isinstance(chat_query, tuple):
-        chat_query = chat_query[0]
-        query = [{'image': chat_query}]
-    else:
-        query = [{'text': _parse_text(chat_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)
-    if 'image' in query[0]:
-        print("Model response:", response)
-        image = tokenizer.draw_bbox_on_latest_picture(response)
-        if image is not None:
-            image_path = save_image(image, uploaded_file_dir)
-            formatted_response = (chat_query, image_path)
-        else:
-            formatted_response = (chat_query, response)
-    else:
-        text_response = response.strip()
-        formatted_response = (chat_query, text_response)
-
-    _chatbot[-1] = formatted_response
-    if task_history:
-        task_history[-1] = formatted_response
-    else:
-        task_history.append(formatted_response)
-
-    return _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) -> 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 add_text(history, task_history, text) -> tuple:
-    if not text.strip():
-        return history, task_history, chatbot
-    if not any(isinstance(item[0], tuple) for item in history):
-        prompt = "Please upload and submit an image to get started."
-        history.append((prompt, None))
-        task_history.append((prompt, None))
-        chatbot.append(prompt)
-        return history, task_history, chatbot
-    task_text = text
-    if len(text) >= 2 and text[-1] in PUNCTUATION and text[-2] not in PUNCTUATION:
-        task_text = text[:-1]
-    history_item = (_parse_text(task_text), None)
-    history.append(history_item)
-    task_history.append(history_item)
-    return history, task_history, chatbot
-
-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: str) -> str:
-    response = response.replace("<ref>", "").replace(r"</ref>", "")
-    response = re.sub(BOX_TAG_PATTERN, "", response)
+def clean_response(response: str) -> str:
+    response = re.sub(r'<ref>(.*?)</ref>(?:<box>.*?</box>)*(?:<quad>.*?</quad>)*', r'\1', response).strip()
     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 _launch_demo(args, model, tokenizer):
-    uploaded_file_dir = os.environ.get("GRADIO_TEMP_DIR") or str(
-        Path(tempfile.gettempdir()) / "gradio"
-    )
+def chat_with_model(image_path=None, text_query=None, history=None):
+    query_elements = []
+    if image_path:
+        query_elements.append({'image': image_path})
+    if text_query:
+        query_elements.append({'text': text_query})
+
+    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
+
+def process_input(text, file):
+    image_path = None
+    if file is not None:
+        image_path = save_image(file, "uploaded_images")
+    response = chat_with_model(image_path=image_path, text_query=text)
+    return response
 
-    with gr.Blocks() as demo:
-        gr.Markdown("""
+with gr.Blocks(theme=ParityError/Anime) as demo:
+    gr.Markdown("""
 # 🙋🏻‍♂️欢迎来到🌟Tonic 的🦆Qwen-VL-Chat🤩Bot！🚀
 # 🙋🏻‍♂️Welcome toTonic's Qwen-VL-Chat Bot! 
 该WebUI基于Qwen-VL-Chat，实现聊天机器人功能。 但我必须解决它的很多问题，也许我也能获得一些荣誉。
@@ -248,51 +59,22 @@ Qwen-VL-Chat 是一种多模式输入模型。 您可以使用此空间来测试
 This WebUI is based on Qwen-VL-Chat, implementing chatbot functionalities. Qwen-VL-Chat is a multimodal input model. You can use this Space to test out the current model [qwen/Qwen-VL-Chat](https://huggingface.co/qwen/Qwen-VL-Chat) You can also use  qwen/Qwen-VL-Chat🚀 by cloning this space.   Simply click here: [Duplicate Space](https://huggingface.co/spaces/Tonic1/VLChat?duplicate=true)
 Join us:  TeamTonic  is always making cool demos! Join our active builder's community on  Discord: [Discord](https://discord.gg/nXx5wbX9) On  Huggingface: [TeamTonic](https://huggingface.co/TeamTonic) & [MultiTransformer](https://huggingface.co/MultiTransformer) On  Github: [Polytonic](https://github.com/tonic-ai) & contribute to   [PolyGPT](https://github.com/tonic-ai/polygpt-alpha)
 """)
-        with gr.Row():
-            with gr.Column(scale=1):
-                chatbot = 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...")
-                    submit_btn = gr.Button("📨Submit")
-                with gr.Row():
-                    file_upload = gr.UploadButton("⤴️Upload Image", file_types=["image"])
-                    submit_file_btn = gr.Button("📩Submit Image")
-                    regen_btn = gr.Button("♻️Regenerate")
-                    empty_bin = gr.Button("🧼Clear History")
-                task_history = gr.State([])
-
-        submit_btn.click(
-            fn=predict,
-            inputs=[chatbot, task_history],
-            outputs=[chatbot]
-        )
-        
-        submit_file_btn.click(
-            fn=handle_image_submission,
-            inputs=[chatbot, task_history, file_upload],
-            outputs=[chatbot, task_history]
-        )
-
-        regen_btn.click(
-            fn=regenerate,
-            inputs=[chatbot, task_history],
-            outputs=[chatbot]
-        )
-
-        empty_bin.click(
-            fn=reset_state,
-            inputs=[task_history],
-            outputs=[task_history],
-        )
-
-        query.submit(
-            fn=add_text,
-            inputs=[chatbot, task_history, query],
-            outputs=[chatbot, task_history, query]
-        )
-
-        gr.Markdown("""
+    with gr.Row():
+        with gr.Column(scale=1):
+            chatbot = 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.File(label="Upload Image")
+                submit_btn = gr.Button("Submit")
+    
+    submit_btn.click(
+        fn=process_input,
+        inputs=[query, file_upload],
+        outputs=chatbot
+    )
+    
+    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,
@@ -300,12 +82,3 @@ including hate speech, violence, pornography, deception, etc. (Note: This demo i
 """)
 
     demo.queue().launch()
-
-
-def main():
-    args = _get_args()
-    model, tokenizer = _load_model_tokenizer(args)
-    _launch_demo(args, model, tokenizer)
-
-if __name__ == '__main__':
-    main()