Update app.py

app.py

@@ -1,237 +1,232 @@
-import gradio as gr
+# Copyright (c) Alibaba Cloud.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
 import os
-import spaces
-from transformers import GemmaTokenizer, AutoModelForCausalLM
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from threading import Thread
-
-# Set an environment variable
-HF_TOKEN = os.environ.get("HF_TOKEN", None)
-
-
-DESCRIPTION = '''
-<div>
-<h1 style="text-align: center;">LLaMA-Mesh</h1>
-<div>
-<a style="display:inline-block" href="https://research.nvidia.com/labs/toronto-ai/LLaMA-Mesh/"><img src='https://img.shields.io/badge/public_website-8A2BE2'></a>
-<a style="display:inline-block; margin-left: .5em" href="https://github.com/nv-tlabs/LLaMA-Mesh"><img src='https://img.shields.io/github/stars/nv-tlabs/LLaMA-Mesh?style=social'/></a>
-</div>
-<p>LLaMA-Mesh: Unifying 3D Mesh Generation with Language Models.<a style="display:inline-block" href="https://research.nvidia.com/labs/toronto-ai/LLaMA-Mesh/">[Project Page]</a> <a style="display:inline-block" href="https://github.com/nv-tlabs/LLaMA-Mesh">[Code]</a></p>
-<p> Notice: (1) This demo supports up to 4096 tokens due to computational limits, while our full model supports 8k tokens. This limitation may result in incomplete generated meshes. To experience the full 8k token context, please run our model locally.</p>
-<p>(2) We only support generating a single mesh per dialog round. To generate another mesh, click the "clear" button and start a new dialog.</p>
-<p>(3) If the LLM refuses to generate a 3D mesh, try adding more explicit instructions to the prompt, such as "create a 3D model of a table <strong>in OBJ format</strong>." A more effective approach is to request the mesh generation at the start of the dialog.</p>
-</div>
-'''
-
-LICENSE = """
-<p/>
----
-Built with Meta Llama 3.1 8B
-"""
-
-PLACEHOLDER = """
-<div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
-   <h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">LLaMA-Mesh</h1>
-   <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.65;">Create 3D meshes by chatting.</p>
-</div>
-"""
-
-
-css = """
-h1 {
-  text-align: center;
-  display: block;
-}
-#duplicate-button {
-  margin: auto;
-  color: white;
-  background: #1565c0;
-  border-radius: 100vh;
-}
-"""
-# Load the tokenizer and model
-model_path = "Zhengyi/LLaMA-Mesh"
-tokenizer = AutoTokenizer.from_pretrained(model_path)
-model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto")
-terminators = [
-    tokenizer.eos_token_id,
-    tokenizer.convert_tokens_to_ids("<|eot_id|>")
-]
-
-
-from trimesh.exchange.gltf import export_glb
-import gradio as gr
-import trimesh
 import numpy as np
+from urllib3.exceptions import HTTPError
+os.system('pip install dashscope  modelscope oss2 -U')
+
+from argparse import ArgumentParser
+from pathlib import Path
+
+import copy
+import gradio as gr
+import oss2
+import os
+import re
+import secrets
 import tempfile
-def apply_gradient_color(mesh_text):
-    """
-    Apply a gradient color to the mesh vertices based on the Y-axis and save as GLB.
-    Args:
-        mesh_text (str): The input mesh in OBJ format as a string.
-    Returns:
-        str: Path to the GLB file with gradient colors applied.
-    """
-    # Load the mesh
-    temp_file =  tempfile.NamedTemporaryFile(suffix=f"", delete=False).name#"temp_mesh.obj"
-    with open(temp_file+".obj", "w") as f:
-        f.write(mesh_text)
-    # return temp_file
-    mesh = trimesh.load_mesh(temp_file+".obj", file_type='obj')
-
-    # Get vertex coordinates
-    vertices = mesh.vertices
-    y_values = vertices[:, 1]  # Y-axis values
-
-    # Normalize Y values to range [0, 1] for color mapping
-    y_normalized = (y_values - y_values.min()) / (y_values.max() - y_values.min())
-
-    # Generate colors: Map normalized Y values to RGB gradient (e.g., blue to red)
-    colors = np.zeros((len(vertices), 4))  # RGBA
-    colors[:, 0] = y_normalized  # Red channel
-    colors[:, 2] = 1 - y_normalized  # Blue channel
-    colors[:, 3] = 1.0  # Alpha channel (fully opaque)
-
-    # Attach colors to mesh vertices
-    mesh.visual.vertex_colors = colors
-
-    # Export to GLB format
-    glb_path = temp_file+".glb"
-    with open(glb_path, "wb") as f:
-        f.write(export_glb(mesh))
-    
-    return glb_path
-
-def visualize_mesh(mesh_text):
-    """
-    Convert the provided 3D mesh text into a visualizable format.
-    This function assumes the input is in OBJ format.
-    """
-    temp_file = "temp_mesh.obj"
-    with open(temp_file, "w") as f:
-        f.write(mesh_text)
-    return temp_file
-
-@spaces.GPU(duration=120)
-def chat_llama3_8b(message: str, 
-              history: list, 
-              temperature: float, 
-              max_new_tokens: int
-             ) -> str:
-    """
-    Generate a streaming response using the llama3-8b model.
-    Args:
-        message (str): The input message.
-        history (list): The conversation history used by ChatInterface.
-        temperature (float): The temperature for generating the response.
-        max_new_tokens (int): The maximum number of new tokens to generate.
-    Returns:
-        str: The generated response.
-    """
-    conversation = []
-    for user, assistant in history:
-        conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
-    conversation.append({"role": "user", "content": message})
-
-    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt").to(model.device)
-    
-    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
-    # print(max_new_tokens)
-    max_new_tokens=4096
-    temperature=0.9
-    generate_kwargs = dict(
-        input_ids= input_ids,
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        temperature=temperature,
-        eos_token_id=terminators,
+import requests
+from http import HTTPStatus
+from dashscope import MultiModalConversation
+import dashscope
+
+API_KEY = os.environ['API_KEY']
+dashscope.api_key = API_KEY
+
+REVISION = 'v1.0.4'
+BOX_TAG_PATTERN = r"<box>([\s\S]*?)</box>"
+PUNCTUATION = "！？。＂＃＄％＆＇（）＊＋，－／：；＜＝＞＠［＼］＾＿｀｛｜｝～｟｠｢｣､、〃》「」『』【】〔〕〖〗〘〙〚〛〜〝〞〟〰〾〿–—‘’‛“”„‟…‧﹏."
+
+
+def _get_args():
+    parser = ArgumentParser()
+    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=7860,
+                        help="Demo server port.")
+    parser.add_argument("--server-name", type=str, default="127.0.0.1",
+                        help="Demo server name.")
+
+    args = parser.parse_args()
+    return args
+
+def _parse_text(text):
+    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
+
+
+def _remove_image_special(text):
+    text = text.replace('<ref>', '').replace('</ref>', '')
+    return re.sub(r'<box>.*?(</box>|$)', '', text)
+
+
+def is_video_file(filename):
+    video_extensions = ['.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv', '.webm', '.mpeg']
+    return any(filename.lower().endswith(ext) for ext in video_extensions)
+
+
+def _launch_demo(args):
+    uploaded_file_dir = os.environ.get("GRADIO_TEMP_DIR") or str(
+        Path(tempfile.gettempdir()) / "gradio"
     )
-    # This will enforce greedy generation (do_sample=False) when the temperature is passed 0, avoiding the crash.             
-    if temperature == 0:
-        generate_kwargs['do_sample'] = False
-        
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
-
-    outputs = []
-    for text in streamer:
-        outputs.append(text)
-        #print(outputs)
-        yield "".join(outputs)
-        
-
-# Gradio block
-chatbot=gr.Chatbot(height=450, placeholder=PLACEHOLDER, label='Gradio ChatInterface')
-
-with gr.Blocks(fill_height=True, css=css) as demo:
-    with gr.Column(): 
-        gr.Markdown(DESCRIPTION)
-        # gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
+
+    def predict(_chatbot, task_history):
+        chat_query = _chatbot[-1][0]
+        query = task_history[-1][0]
+        if len(chat_query) == 0:
+            _chatbot.pop()
+            task_history.pop()
+            return _chatbot
+        print("User: " + _parse_text(query))
+        history_cp = copy.deepcopy(task_history)
+        full_response = ""
+        messages = []
+        content = []
+        for q, a in history_cp:
+            if isinstance(q, (tuple, list)):
+                if is_video_file(q[0]):
+                    content.append({'video': f'file://{q[0]}'})
+                else:
+                    content.append({'image': f'file://{q[0]}'})
+            else:
+                content.append({'text': q})
+                messages.append({'role': 'user', 'content': content})
+                messages.append({'role': 'assistant', 'content': [{'text': a}]})
+                content = []
+        messages.pop()
+        responses = MultiModalConversation.call(
+            model='qwen2.5-vl-32b-instruct', messages=messages, stream=True,
+        )
+        for response in responses:
+            if not response.status_code == HTTPStatus.OK:
+                raise HTTPError(f'response.code: {response.code}\nresponse.message: {response.message}')
+            response = response.output.choices[0].message.content
+            response_text = []
+            for ele in response:
+                if 'text' in ele:
+                    response_text.append(ele['text'])
+                elif 'box' in ele:
+                    response_text.append(ele['box'])
+            response_text = ''.join(response_text)
+            _chatbot[-1] = (_parse_text(chat_query), _remove_image_special(response_text))
+            yield _chatbot
+
+        if len(response) > 1:
+            result_image = response[-1]['result_image']
+            resp = requests.get(result_image)
+            os.makedirs(uploaded_file_dir, exist_ok=True)
+            name = f"tmp{secrets.token_hex(20)}.jpg"
+            filename = os.path.join(uploaded_file_dir, name)
+            with open(filename, 'wb') as f:
+                f.write(resp.content)
+            response = ''.join(r['box'] if 'box' in r else r['text'] for r in response[:-1])
+            _chatbot.append((None, (filename,)))
+        else:
+            response = response[0]['text']
+            _chatbot[-1] = (_parse_text(chat_query), response)
+        full_response = _parse_text(response)
+
+        task_history[-1] = (query, full_response)
+        print("Qwen2.5-VL-Chat: " + _parse_text(full_response))
+        yield _chatbot
+
+
+    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))
+        _chatbot_gen = predict(_chatbot, task_history)
+        for _chatbot in _chatbot_gen:
+            yield _chatbot
+
+    def add_text(history, task_history, text):
+        task_text = text
+        history = history if history is not None else []
+        task_history = task_history if task_history is not None else []
+        history = history + [(_parse_text(text), None)]
+        task_history = task_history + [(task_text, None)]
+        return history, task_history, ""
+
+    def add_file(history, task_history, file):
+        history = history if history is not None else []
+        task_history = task_history if task_history is not None else []
+        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 reset_state(task_history):
+        task_history.clear()
+        return []
+
+    with gr.Blocks() as demo:
+        gr.Markdown("""<center><font size=3> Qwen2.5-VL-32B-Instruct Demo </center>""")
+
+        chatbot = gr.Chatbot(label='Qwen2.5-VL-32B-Instruct', elem_classes="control-height", height=500)
+        query = gr.Textbox(lines=2, label='Input')
+        task_history = gr.State([])
+
         with gr.Row():
-            with gr.Column(scale=3):    
-                gr.ChatInterface(
-                    fn=chat_llama3_8b,
-                    chatbot=chatbot,
-                    fill_height=True,
-                    additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False, render=False),
-                    additional_inputs=[
-                        gr.Slider(minimum=0,
-                                maximum=1, 
-                                step=0.1,
-                                value=0.9, 
-                                label="Temperature", 
-                                interactive = False,
-                                render=False),
-                        gr.Slider(minimum=128, 
-                                maximum=4096,
-                                step=1,
-                                value=4096, 
-                                label="Max new tokens", 
-                                interactive = False,
-                                render=False),
-                        ],
-                    examples=[
-                        ['Create a 3D model of a wooden hammer'],
-                        ['Create a 3D model of a pyramid in obj format'],
-                        ['Create a 3D model of a cabinet.'],
-                        ['Create a low poly 3D model of a coffe cup'],
-                        ['Create a 3D model of a table.'],
-                        ["Create a low poly 3D model of a tree."],
-                        ['Write a python code for sorting.'],
-                        ['How to setup a human base on Mars? Give short answer.'],
-                        ['Explain theory of relativity to me like I’m 8 years old.'],
-                        ['What is 9,000 * 9,000?'],
-                        ['Create a 3D model of a soda can.'],
-                        ['Create a 3D model of a sword.'],
-                        ['Create a 3D model of a wooden barrel'],
-                        ['Create a 3D model of a chair.']
-                        ],
-                    cache_examples=False,
-                                )
-                gr.Markdown(LICENSE)
-        
-            with gr.Column(scale=2): 
-                output_model = gr.Model3D(
-                            label="3D Mesh Visualization",
-                            interactive=False,
-                        )
-                gr.Markdown("You can copy the generated 3d objects in the left and paste in the textbox below. Put the button and you will see the visualization of the 3D mesh.")
-                
-                # Add the text box for 3D mesh input and button
-                mesh_input = gr.Textbox(
-                    label="3D Mesh Input",
-                    placeholder="Paste your 3D mesh in OBJ format here...",
-                    lines=5,
-                )
-                visualize_button = gr.Button("Visualize 3D Mesh")
-                
-                # Link the button to the visualization function
-                visualize_button.click(
-                    fn=apply_gradient_color,
-                    inputs=[mesh_input],
-                    outputs=[output_model]
-                )
-          
-if __name__ == "__main__":
-    demo.launch()
+            addfile_btn = gr.UploadButton("📁 Upload (上传文件)", file_types=["image", "video"])
+            submit_btn = gr.Button("🚀 Submit (发送)")
+            regen_btn = gr.Button("🤔️ Regenerate (重试)")
+            empty_bin = gr.Button("🧹 Clear History (清除历史)")
+
+        submit_btn.click(add_text, [chatbot, task_history, query], [chatbot, task_history]).then(
+            predict, [chatbot, task_history], [chatbot], show_progress=True
+        )
+        submit_btn.click(reset_user_input, [], [query])
+        empty_bin.click(reset_state, [task_history], [chatbot], show_progress=True)
+        regen_btn.click(regenerate, [chatbot, task_history], [chatbot], show_progress=True)
+        addfile_btn.upload(add_file, [chatbot, task_history, addfile_btn], [chatbot, task_history], show_progress=True)
+
+
+    demo.queue(default_concurrency_limit=40).launch(
+        share=args.share,
+        # inbrowser=args.inbrowser,
+        # server_port=args.server_port,
+        # server_name=args.server_name,
+    )
+
+
+def main():
+    args = _get_args()
+    _launch_demo(args)
+
+
+if __name__ == '__main__':
+    main()