update

app.py

@@ -7,8 +7,14 @@ import re
 import argparse
 import copy
 import spaces
+import fitz
+import subprocess
+import tempfile
+import os
+import time
 
-MODEL_PATH = "/model/glm-4v-9b-0529"
+MODEL_PATH = "THUDM/GLM-4.1V-9B-Thinking"
+stop_generation = False
 
 
 class GLM4VModel:
@@ -32,6 +38,26 @@ class GLM4VModel:
     def _wrap_text(self, t):
         return [{"type": "text", "text": t}]
 
+    def _pdf_to_imgs(self, pdf_path):
+        doc = fitz.open(pdf_path)
+        imgs = []
+        for i in range(doc.page_count):
+            pix = doc.load_page(i).get_pixmap(dpi=180)
+            img_p = os.path.join(tempfile.gettempdir(), f"{Path(pdf_path).stem}_{i}.png")
+            pix.save(img_p)
+            imgs.append(img_p)
+        doc.close()
+        return imgs
+
+    def _ppt_to_imgs(self, ppt_path):
+        tmp = tempfile.mkdtemp()
+        subprocess.run(
+            ["libreoffice", "--headless", "--convert-to", "pdf", "--outdir", tmp, ppt_path],
+            check=True,
+        )
+        pdf_path = os.path.join(tmp, Path(ppt_path).stem + ".pdf")
+        return self._pdf_to_imgs(pdf_path)
+
     def _files_to_content(self, media):
         out = []
         for f in media or []:
@@ -40,94 +66,70 @@ class GLM4VModel:
                 out.append({"type": "video", "url": f.name})
             elif ext in [".jpg", ".jpeg", ".png", ".gif", ".bmp", ".tiff", ".webp"]:
                 out.append({"type": "image", "url": f.name})
+            elif ext in [".ppt", ".pptx"]:
+                for p in self._ppt_to_imgs(f.name):
+                    out.append({"type": "image", "url": p})
+            elif ext == ".pdf":
+                for p in self._pdf_to_imgs(f.name):
+                    out.append({"type": "image", "url": p})
         return out
 
-    # -----------------------------------------------------------
-    # 🖼️  Output formatting
-    # -----------------------------------------------------------
-    def _format_output(self, txt):
-        """Called once完整生成结束时"""
-        think_pat, ans_pat = r"<think>(.*?)</think>", r"<answer>(.*?)</answer>"
-        think = re.findall(think_pat, txt, re.DOTALL)
-        ans = re.findall(ans_pat, txt, re.DOTALL)
-        html = ""
-        if think:
-            html += (
-                    "<details open><summary style='cursor:pointer;font-weight:bold;color:#bbbbbb;'>💭 Thinking Process</summary>"
-                    "<div style='color:#cccccc;line-height:1.4;'>"
-                    + think[0].strip()
-                    + "</div></details><br>"
-            )
-        body = ans[0] if ans else re.sub(think_pat, "", txt, flags=re.DOTALL)
-        html += f"<div style='color:#ffffff;'>{body.strip()}</div>"
-        return html
-
     def _stream_fragment(self, buf: str) -> str:
         think_html = ""
         if "<think>" in buf:
             if "</think>" in buf:
-                think_content = re.search(r"<think>(.*?)</think>", buf, re.DOTALL)
-                if think_content:
+                seg = re.search(r"<think>(.*?)</think>", buf, re.DOTALL)
+                if seg:
                     think_html = (
-                            "<details open><summary style='cursor:pointer;font-weight:bold;color:#bbbbbb;'>💭 Thinking Process</summary>"
-                            "<div style='color:#cccccc;line-height:1.4;'>"
-                            + think_content.group(1).strip()
-                            + "</div></details><br>"
+                        "<details open><summary style='cursor:pointer;font-weight:bold;color:#bbbbbb;'>💭 思考过程</summary>"
+                        "<div style='color:#cccccc;line-height:1.4;padding:10px;border-left:3px solid #666;margin:5px 0;background-color:rgba(128,128,128,0.1);'>"
+                        + seg.group(1).strip().replace("\n", "<br>")
+                        + "</div></details>"
                     )
             else:
-                partial = buf.split("<think>", 1)[1]
+                part = buf.split("<think>", 1)[1]
                 think_html = (
-                        "<details open><summary style='cursor:pointer;font-weight:bold;color:#bbbbbb;'>💭 Thinking Process</summary>"
-                        "<div style='color:#cccccc;line-height:1.4;'>" + partial
+                    "<details open><summary style='cursor:pointer;font-weight:bold;color:#bbbbbb;'>💭 思考过程</summary>"
+                    "<div style='color:#cccccc;line-height:1.4;padding:10px;border-left:3px solid #666;margin:5px 0;background-color:rgba(128,128,128,0.1);'>"
+                    + part.replace("\n", "<br>")
+                    + "</div></details>"
                 )
 
         answer_html = ""
         if "<answer>" in buf:
             if "</answer>" in buf:
-                ans_content = re.search(r"<answer>(.*?)</answer>", buf, re.DOTALL)
-                if ans_content:
-                    answer_html = (
-                            "<div style='color:#ffffff;'>" + ans_content.group(1).strip() + "</div>"
-                    )
+                seg = re.search(r"<answer>(.*?)</answer>", buf, re.DOTALL)
+                if seg:
+                    answer_html = seg.group(1).strip()
             else:
-                partial = buf.split("<answer>", 1)[1]
-                answer_html = "<div style='color:#ffffff;'>" + partial
+                answer_html = buf.split("<answer>", 1)[1]
 
         if not think_html and not answer_html:
             return self._strip_html(buf)
-
         return think_html + answer_html
 
-    def _build_messages(self, hist, sys_prompt):
+    def _build_messages(self, raw_hist, sys_prompt):
         msgs = []
 
         if sys_prompt.strip():
-            msgs.append({
-                "role": "system",
-                "content": [{"type": "text", "text": sys_prompt.strip()}]
-            })
+            msgs.append({"role": "system", "content": [{"type": "text", "text": sys_prompt.strip()}]})
 
-        for h in hist:
+        for h in raw_hist:
             if h["role"] == "user":
-                payload = h.get("file_info") or self._wrap_text(
-                    self._strip_html(h["content"])
-                )
-                msgs.append({"role": "user", "content": payload})
-
+                msgs.append({"role": "user", "content": h["content"]})
             else:
                 raw = h["content"]
                 raw = re.sub(r"<think>.*?</think>", "", raw, flags=re.DOTALL)
                 raw = re.sub(r"<details.*?</details>", "", raw, flags=re.DOTALL)
                 clean = self._strip_html(raw).strip()
-
                 msgs.append({"role": "assistant", "content": self._wrap_text(clean)})
-
         return msgs
 
     @spaces.GPU(duration=240)
-    def stream_generate(self, hist, sys_prompt):
-        msgs = self._build_messages(hist, sys_prompt)
-        print(msgs)
+    def stream_generate(self, raw_hist, sys_prompt):
+        global stop_generation
+        stop_generation = False
+        msgs = self._build_messages(raw_hist, sys_prompt)
         inputs = self.processor.apply_chat_template(
             msgs,
             tokenize=True,
@@ -137,9 +139,7 @@ class GLM4VModel:
             padding=True,
         ).to(self.device)
 
-        streamer = TextIteratorStreamer(
-            self.processor.tokenizer, skip_prompt=True, skip_special_tokens=False
-        )
+        streamer = TextIteratorStreamer(self.processor.tokenizer, skip_prompt=True, skip_special_tokens=False)
         gen_args = dict(
             inputs,
             max_new_tokens=8192,
@@ -150,13 +150,46 @@ class GLM4VModel:
             top_p=1e-5,
             streamer=streamer,
         )
-        threading.Thread(target=self.model.generate, kwargs=gen_args).start()
+
+        generation_thread = threading.Thread(target=self.model.generate, kwargs=gen_args)
+        generation_thread.start()
 
         buf = ""
         for tok in streamer:
+            if stop_generation:
+                break
             buf += tok
             yield self._stream_fragment(buf)
-        yield self._format_output(buf)
+
+        generation_thread.join()
+
+
+def format_display_content(content):
+    if isinstance(content, list):
+        text_parts = []
+        file_count = 0
+        for item in content:
+            if item["type"] == "text":
+                text_parts.append(item["text"])
+            else:
+                file_count += 1
+
+        display_text = " ".join(text_parts)
+        if file_count > 0:
+            return f"[{file_count} file(s) uploaded]\n{display_text}"
+        return display_text
+    return content
+
+
+def create_display_history(raw_hist):
+    display_hist = []
+    for h in raw_hist:
+        if h["role"] == "user":
+            display_content = format_display_content(h["content"])
+            display_hist.append({"role": "user", "content": display_content})
+        else:
+            display_hist.append({"role": "assistant", "content": h["content"]})
+    return display_hist
 
 
 glm4v = GLM4VModel()
@@ -164,27 +197,35 @@ glm4v.load()
 
 
 def check_files(files):
-    vids = imgs = 0
+    vids = imgs = ppts = pdfs = 0
     for f in files or []:
         ext = Path(f.name).suffix.lower()
         if ext in [".mp4", ".avi", ".mkv", ".mov", ".wmv", ".flv", ".webm", ".mpeg", ".m4v"]:
             vids += 1
         elif ext in [".jpg", ".jpeg", ".png", ".gif", ".bmp", ".tiff", ".webp"]:
             imgs += 1
-    if vids > 1:
-        return False, "Only 1 video allowed"
+        elif ext in [".ppt", ".pptx"]:
+            ppts += 1
+        elif ext == ".pdf":
+            pdfs += 1
+    if vids > 1 or ppts > 1 or pdfs > 1:
+        return False, "Only one video or one PPT or one PDF allowed"
     if imgs > 10:
-        return False, "Max 10 images"
-    if vids and imgs:
-        return False, "Cannot mix video and images"
+        return False, "Maximum 10 images allowed"
+    if (ppts or pdfs) and (vids or imgs) or (vids and imgs):
+        return False, "Cannot mix documents, videos, and images"
     return True, ""
 
 
-def chat(files, msg, hist, sys_prompt):
+def chat(files, msg, raw_hist, sys_prompt):
+    global stop_generation
+    stop_generation = False
+
     ok, err = check_files(files)
     if not ok:
-        hist.append({"role": "assistant", "content": err})
-        yield copy.deepcopy(hist), None, ""
+        raw_hist.append({"role": "assistant", "content": err})
+        display_hist = create_display_history(raw_hist)
+        yield display_hist, copy.deepcopy(raw_hist), None, ""
         return
 
     payload = glm4v._files_to_content(files) if files else None
@@ -194,24 +235,33 @@ def chat(files, msg, hist, sys_prompt):
         else:
             payload.append({"type": "text", "text": msg.strip()})
 
-    display = f"[{len(files)} file(s) uploaded]\n{msg}" if files else msg
-    user_rec = {"role": "user", "content": display}
-    if payload:
-        user_rec["file_info"] = payload
-    hist.append(user_rec)
+    user_rec = {"role": "user", "content": payload if payload else msg.strip()}
+    if raw_hist is None:
+        raw_hist = []
+    raw_hist.append(user_rec)
 
     place = {"role": "assistant", "content": ""}
-    hist.append(place)
-    yield copy.deepcopy(hist), None, ""
+    raw_hist.append(place)
+
+    display_hist = create_display_history(raw_hist)
+    yield display_hist, copy.deepcopy(raw_hist), None, ""
 
-    for chunk in glm4v.stream_generate(hist[:-1], sys_prompt):
+    for chunk in glm4v.stream_generate(raw_hist[:-1], sys_prompt):
+        if stop_generation:
+            break
         place["content"] = chunk
-        yield copy.deepcopy(hist), None, ""
-    yield copy.deepcopy(hist), None, ""
+        display_hist = create_display_history(raw_hist)
+        yield display_hist, copy.deepcopy(raw_hist), None, ""
+
+    display_hist = create_display_history(raw_hist)
+    yield display_hist, copy.deepcopy(raw_hist), None, ""
 
 
 def reset():
-    return [], None, ""
+    global stop_generation
+    stop_generation = True
+    time.sleep(0.1)
+    return [], [], None, ""
 
 
 css = """.chatbot-container .message-wrap .message{font-size:14px!important}
@@ -222,14 +272,16 @@ demo = gr.Blocks(title="GLM-4.1V Chat", theme=gr.themes.Soft(), css=css)
 with demo:
     gr.Markdown("""
                <div style="text-align: center; font-size: 32px; font-weight: bold; margin-bottom: 20px;">
-                   GLM-4.1V-9B Gradio Space🤗
+                   GLM-4.1V-9B-Thinking Gradio Space🤗
                 </div>
                <div style="text-align: center;">
                <a href="https://huggingface.co/THUDM/GLM-4.1V-9B-Thinking">🤗 Model Hub</a> | 
-               <a href="https://github.com/THUDM/CogVLM">🌐 Github</a> |
-               <a href="https://arxiv.org/abs/">📜 arxiv</a>
+               <a href="https://github.com/THUDM/GLM-4.1V-Thinking">🌐 Github</a> 
                 </div>
                 """)
+
+    raw_history = gr.State([])
+
     with gr.Row():
         with gr.Column(scale=7):
             chatbox = gr.Chatbot(
@@ -246,27 +298,18 @@ with demo:
             up = gr.File(
                 label="📁 Upload",
                 file_count="multiple",
-                file_types=["image", "video"],
+                file_types=["file"],
                 type="filepath",
             )
-            gr.Markdown("""
-            <span style="color:red">Please upload the Bay image before entering text.</span>
-            """)
+            gr.Markdown("Supports images / videos / PPT / PDF")
+            gr.Markdown(
+                "The maximum supported input is 10 images or 1 video/PPT/PDF. During the conversation, video and images cannot be present at the same time."
+            )
             sys = gr.Textbox(label="⚙️ System Prompt", lines=6)
 
-    send.click(chat, inputs=[up, textbox, chatbox, sys], outputs=[chatbox, up, textbox])
-    textbox.submit(chat, inputs=[up, textbox, chatbox, sys], outputs=[chatbox, up, textbox])
-    clear.click(reset, outputs=[chatbox, up, textbox])
+    send.click(chat, inputs=[up, textbox, raw_history, sys], outputs=[chatbox, raw_history, up, textbox])
+    textbox.submit(chat, inputs=[up, textbox, raw_history, sys], outputs=[chatbox, raw_history, up, textbox])
+    clear.click(reset, outputs=[chatbox, raw_history, up, textbox])
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--port", type=int, default=8000)
-    parser.add_argument("--host", type=str, default="0.0.0.0")
-    parser.add_argument("--share", action="store_true")
-    args = parser.parse_args()
-
-    demo.launch(
-        server_port=args.port,
-        server_name=args.host,
-        share=args.share,
-    )
+    demo.launch()