Update app.py

app.py

@@ -4,27 +4,119 @@ from threading import Thread
 import gradio as gr
 from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor,TextIteratorStreamer,AutoTokenizer
 from qwen_vl_utils import process_vision_info
-
-# 3D mesh dependencies
 import trimesh
 from trimesh.exchange.gltf import export_glb
 import numpy as np
 import tempfile
 
+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()
+
+    messages = _transform_messages(messages)
+    text = processor.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True)
+    image_inputs, video_inputs = process_vision_info(messages)
+    inputs = processor(text=[text], images=image_inputs,
+                        videos=video_inputs, padding=True, return_tensors='pt')
+    inputs = inputs.to(model.device)
+
+    streamer = TextIteratorStreamer(
+        tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
+
+    gen_kwargs = {'max_new_tokens': 512, 'streamer': streamer, **inputs}
+
+    thread = Thread(target=model.generate, kwargs=gen_kwargs)
+    thread.start()
+
+    #for new_text in streamer:
+    #    yield new_text
+
+    buffer = []
+    for chunk in streamer:
+        buffer.append(chunk)
+        yield "".join(buffer)
+
+
+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 []
+
+def _transform_messages(original_messages):
+    transformed_messages = []
+    for message in original_messages:
+        new_content = []
+        for item in message['content']:
+            if 'image' in item:
+                new_item = {'type': 'image', 'image': item['image']}
+            elif 'text' in item:
+                new_item = {'type': 'text', 'text': item['text']}
+            elif 'video' in item:
+                new_item = {'type': 'video', 'video': item['video']}
+            else:
+                continue
+            new_content.append(new_item)
+
+        new_message = {'role': message['role'], 'content': new_content}
+        transformed_messages.append(new_message)
+
+    return transformed_messages
 
-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>
-'''
 # --------- Configuration & Model Loading ---------
 MODEL_DIR = "Qwen/Qwen2.5-VL-3B-Instruct"
 # Load processor, tokenizer, model for Qwen2.5-VL
@@ -35,90 +127,43 @@ model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     trust_remote_code=True
 )
 processor = AutoProcessor.from_pretrained(MODEL_DIR)
-terminators = [tokenizer.eos_token_id]
+tokenizer = processor.tokenizer
+#terminators = [tokenizer.eos_token_id]
 
-def chat_qwen_vl(message: str, history: list, temperature: float = 0.1, max_new_tokens: int = 1024):
-    # —— 原有多模态输入构造 —— #
+def chat_qwen_vl(messages: str, history: list, temperature: float = 0.1, max_new_tokens: int = 1024):
     messages = [
             {
                 "role": "user",
                 "content": [
-                    {"type": "text", "text": message},
+                    {"type": "text", "text": messages},
                 ],
             }
         ]
+    messages = _transform_messages(messages)
     text = processor.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
-    print(text)
+        messages, tokenize=False, add_generation_prompt=True)
     image_inputs, video_inputs = process_vision_info(messages)
-    input_ids = processor(
-        text=[text],
-        images=image_inputs,
-        videos=video_inputs,
-        padding=False,
-        return_tensors="pt"
-    ).to(model.device)
-
-    # 2. 把 streamer 和生成参数一起传给 model.generate
+    inputs = processor(text=[text], images=image_inputs,
+                        videos=video_inputs, padding=True, return_tensors='pt')
+    inputs = inputs.to(model.device)
+
     streamer = TextIteratorStreamer(
-        processor,
-        timeout=100.0,
-        skip_prompt=True,
-        skip_special_tokens=True
-    )
-    #print(input_ids)
-    generate_kwargs = dict(
-        input_ids= input_ids["input_ids"],
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        temperature=temperature,
-        eos_token_id=terminators,
-        top_k=1024,
-        top_p=0.1
-    )
+        tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
 
-    # 4. 后台线程启动生成
-    Thread(target=model.generate, kwargs=generate_kwargs).start()
+    gen_kwargs = {'max_new_tokens': 512, 'streamer': streamer, **inputs}
+
+    thread = Thread(target=model.generate, kwargs=gen_kwargs)
+    thread.start()
+
+    #for new_text in streamer:
+    #    yield new_text
 
-    # 5. 主线程读取 streamer 并 yield
     buffer = []
     for chunk in streamer:
-        print(chunk)
         buffer.append(chunk)
-        # 每次新到一个片段，就拼接并返回给前端
         yield "".join(buffer)
-        
-        
-# --------- 3D Mesh Coloring Function ---------
-def apply_gradient_color(mesh_text: str) -> str:
-    """
-    Apply a Y-axis-based gradient RGBA color to OBJ mesh text and export as GLB.
-    """
-    # Write OBJ to temp file
-    tmp = tempfile.NamedTemporaryFile(suffix=".obj", delete=False)
-    tmp.write(mesh_text.encode('utf-8'))
-    tmp.flush()
-    tmp.close()
-
-    mesh = trimesh.load_mesh(tmp.name, file_type='obj')
-    vertices = mesh.vertices
-    ys = vertices[:, 1]
-    y_norm = (ys - ys.min()) / (ys.max() - ys.min())
-
-    colors = np.zeros((len(vertices), 4))
-    colors[:, 0] = y_norm
-    colors[:, 2] = 1 - y_norm
-    colors[:, 3] = 1.0
-    mesh.visual.vertex_colors = colors
-
-    glb_path = tmp.name.replace('.obj', '.glb')
-    with open(glb_path, 'wb') as f:
-        f.write(export_glb(mesh))
-    return glb_path
-
-# --------- Gradio Interface ---------
+
+
 css = """
 h1 { text-align: center; }
 """
@@ -128,53 +173,27 @@ PLACEHOLDER = (
     "<p style='font-size:18px;opacity:0.65;'>Ask anything or generate images!</p></div>"
 )
 
-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")
-        with gr.Row():
-            with gr.Column(scale=3):    
-                gr.ChatInterface(
-                    fn=chat_qwen_vl,
-                    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,
-                                )
+with gr.Blocks() as demo:
+    gr.Markdown("""<center><font size=3> ShapeLLM-7B Demo </center>""")
+
+    chatbot = gr.Chatbot(label='ShapeLLM-4o', elem_classes="control-height", height=500)
+    query = gr.Textbox(lines=2, label='Input')
+    task_history = gr.State([])
+
+    with gr.Row():
+        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)
+    
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()