1

app.py

@@ -1,17 +1,17 @@
-import gradio as gr
-import torch
-from transformers import AutoProcessor, Glm4vForConditionalGeneration, TextIteratorStreamer
-from pathlib import Path
-import threading
-import re
 import copy
-import spaces
-import fitz
+import os
+import re
 import subprocess
 import tempfile
-import os
+import threading
+from pathlib import Path
+import fitz
+import gradio as gr
 import time
 import html
+import torch
+from transformers import AutoProcessor, Glm4vForConditionalGeneration, TextIteratorStreamer
+import spaces
 
 MODEL_PATH = "THUDM/GLM-4.1V-9B-Thinking"
 stop_generation = False
@@ -20,142 +20,142 @@ processor = AutoProcessor.from_pretrained(MODEL_PATH, use_fast=True)
 model = Glm4vForConditionalGeneration.from_pretrained(
     MODEL_PATH,
     torch_dtype=torch.bfloat16,
-    device_map="auto",
-    attn_implementation="sdpa",
+    device_map="auto"
 )
 
 
-def _strip_html(text):
-    return re.sub(r"<[^>]+>", "", text).strip()
-
+class GLM4VModel:
+    def _strip_html(self, text: str) -> str:
+        return re.sub(r"<[^>]+>", "", text).strip()
 
-def _wrap_text(text):
-    return [{"type": "text", "text": text}]
+    def _wrap_text(self, text: str):
+        return [{"type": "text", "text": text}]
 
+    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 _pdf_to_imgs(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(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 _pdf_to_imgs(pdf_path)
+    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 []:
+            ext = Path(f.name).suffix.lower()
+            if ext in [".mp4", ".avi", ".mkv", ".mov", ".wmv", ".flv", ".webm", ".mpeg", ".m4v"]:
+                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
 
-def _files_to_content(media):
-    out = []
-    for f in media or []:
-        ext = Path(f.name).suffix.lower()
-        if ext in [".mp4", ".avi", ".mkv", ".mov", ".wmv", ".flv", ".webm", ".mpeg", ".m4v"]:
-            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 _ppt_to_imgs(f.name):
-                out.append({"type": "image", "url": p})
-        elif ext == ".pdf":
-            for p in _pdf_to_imgs(f.name):
-                out.append({"type": "image", "image_url": p})
-    return out
-
-
-def _stream_fragment(buf, skip_think=False):
-    think_html = ""
-    if "<think>" in buf and not skip_think:
-        if "</think>" in buf:
-            seg = re.search(r"<think>(.*?)</think>", buf, re.DOTALL)
-            if seg:
-                think_content = seg.group(1).strip().replace("\\n", "\n").replace("\n", "<br>")
+    def _stream_fragment(self, buf: str, skip_think: bool = False):
+        think_html = ""
+        if "<think>" in buf and not skip_think:
+            if "</think>" in buf:
+                seg = re.search(r"<think>(.*?)</think>", buf, re.DOTALL)
+                if seg:
+                    think_content = seg.group(1).strip().replace("\\n", "\n").replace("\n", "<br>")
+                    think_html = (
+                            "<details open><summary style='cursor:pointer;font-weight:bold;color:#007acc;'>💭 Thinking</summary>"
+                            "<div style='color:#555555;line-height:1.6;padding:15px;border-left:4px solid #007acc;margin:10px 0;background-color:#f0f7ff;border-radius:4px;'>"
+                            + think_content
+                            + "</div></details>"
+                    )
+            else:
+                part = buf.split("<think>", 1)[1]
+                think_content = part.replace("\\n", "\n").replace("\n", "<br>")
                 think_html = (
                         "<details open><summary style='cursor:pointer;font-weight:bold;color:#007acc;'>💭 Thinking</summary>"
                         "<div style='color:#555555;line-height:1.6;padding:15px;border-left:4px solid #007acc;margin:10px 0;background-color:#f0f7ff;border-radius:4px;'>"
-                        + think_content + "</div></details>"
+                        + think_content
+                        + "</div></details>"
                 )
-        else:
-            part = buf.split("<think>", 1)[1]
-            think_content = part.replace("\\n", "\n").replace("\n", "<br>")
-            think_html = (
-                    "<details open><summary style='cursor:pointer;font-weight:bold;color:#007acc;'>💭 Thinking</summary>"
-                    "<div style='color:#555555;line-height:1.6;padding:15px;border-left:4px solid #007acc;margin:10px 0;background-color:#f0f7ff;border-radius:4px;'>"
-                    + think_content + "</div></details>"
-            )
-    answer_html = ""
-    if "<answer>" in buf:
-        if "</answer>" in buf:
-            seg = re.search(r"<answer>(.*?)</answer>", buf, re.DOTALL)
-            if seg:
-                answer_html = seg.group(1).strip()
-        else:
-            answer_html = buf.split("<answer>", 1)[1]
-    if answer_html:
-        answer_html_raw = answer_html.replace("\\n", "\n")
-        escaped = html.escape(answer_html_raw)
-        answer_html = f"<pre class='code-block'><code class='language-html'>{escaped}</code></pre>"
-    if not think_html and not answer_html:
-        return _strip_html(buf)
-    return think_html + answer_html
-
-
-def _build_messages(raw_hist, sys_prompt):
-    msgs = []
-    if sys_prompt.strip():
-        msgs.append({"role": "system", "content": [{"type": "text", "text": sys_prompt.strip()}]})
-    for h in raw_hist:
-        if h["role"] == "user":
-            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)
-            msgs.append({"role": "assistant", "content": _wrap_text(_strip_html(raw).strip())})
-    return msgs
 
+        answer_html = ""
+        if "<answer>" in buf:
+            if "</answer>" in buf:
+                seg = re.search(r"<answer>(.*?)</answer>", buf, re.DOTALL)
+                if seg:
+                    answer_html = seg.group(1).strip()
+            else:
+                answer_html = buf.split("<answer>", 1)[1]
 
-@spaces.GPU()
-def stream_generate(raw_hist, sys_prompt):
-    global stop_generation
-    stop_generation = False
-    msgs = _build_messages(raw_hist, sys_prompt)
-    inputs = processor.apply_chat_template(
-        msgs,
-        tokenize=True,
-        add_generation_prompt=True,
-        return_dict=True,
-        return_tensors="pt",
-        padding=True,
-    ).to(model.device)
-    streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=False)
-    gen_kwargs = dict(
-        inputs,
-        max_new_tokens=8192,
-        repetition_penalty=1.1,
-        do_sample=True,
-        top_k=2,
-        temperature=0.01,
-        streamer=streamer,
-    )
-    thread = threading.Thread(target=model.generate, kwargs=gen_kwargs)
-    thread.start()
-    buf = ""
-    for tok in streamer:
-        if stop_generation:
-            break
-        buf += tok
-        yield _stream_fragment(buf)
-    thread.join()
+        if answer_html:
+            answer_html_raw = answer_html.replace("\\n", "\n")
+            if '<' in answer_html_raw and '>' in answer_html_raw:
+                escaped = html.escape(answer_html_raw)
+                answer_html = f"<pre class='code-block'><code class='language-html'>{escaped}</code></pre>"
+            else:
+                answer_html = f"<div style='margin:0.5em 0; white-space: pre-wrap; line-height:1.6;'>{html.escape(answer_html_raw)}</div>"
+
+        if not think_html and not answer_html:
+            return self._strip_html(buf)
+        return think_html + answer_html
+
+    def _build_messages(self, raw_hist, sys_prompt):
+        msgs = []
+        if sys_prompt.strip():
+            msgs.append({"role": "system", "content": [{"type": "text", "text": sys_prompt.strip()}]})
+        for h in raw_hist:
+            if h["role"] == "user":
+                msgs.append({"role": "user", "content": h["content"]})
+            else:
+                raw = re.sub(r"<think>.*?</think>", "", h["content"], flags=re.DOTALL)
+                raw = re.sub(r"<details.*?</details>", "", raw, flags=re.DOTALL)
+                msgs.append({"role": "assistant", "content": self._wrap_text(self._strip_html(raw).strip())})
+        return msgs
+
+    @spaces.GPU(duration=120)
+    def stream_generate(self, raw_hist, sys_prompt: str):
+        global stop_generation
+        stop_generation = False
+        msgs = self._build_messages(raw_hist, sys_prompt)
+        inputs = processor.apply_chat_template(
+            msgs,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_dict=True,
+            return_tensors="pt",
+            padding=True,
+        ).to(model.device)
+        streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=False)
+        gen_kwargs = dict(
+            inputs,
+            max_new_tokens=8192,
+            repetition_penalty=1.1,
+            do_sample=True,
+            top_k=2,
+            temperature=0.01,
+            streamer=streamer,
+        )
+        thread = threading.Thread(target=model.generate, kwargs=gen_kwargs)
+        thread.start()
+        buf = ""
+        for tok in streamer:
+            if stop_generation:
+                break
+            buf += tok
+            yield self._stream_fragment(buf)
+        thread.join()
 
 
 def format_display_content(content):
@@ -185,6 +185,9 @@ def create_display_history(raw_hist):
     return display_hist
 
 
+glm4v = GLM4VModel()
+
+
 def check_files(files):
     vids = imgs = ppts = pdfs = 0
     for f in files or []:
@@ -215,10 +218,10 @@ def chat(files, msg, raw_hist, sys_prompt):
         display_hist = create_display_history(raw_hist)
         yield display_hist, copy.deepcopy(raw_hist), None, ""
         return
-    payload = _files_to_content(files) if files else None
+    payload = glm4v._files_to_content(files) if files else None
     if msg.strip():
         if payload is None:
-            payload = _wrap_text(msg.strip())
+            payload = glm4v._wrap_text(msg.strip())
         else:
             payload.append({"type": "text", "text": msg.strip()})
     user_rec = {"role": "user", "content": payload if payload else msg.strip()}
@@ -229,7 +232,7 @@ def chat(files, msg, raw_hist, sys_prompt):
     raw_hist.append(place)
     display_hist = create_display_history(raw_hist)
     yield display_hist, copy.deepcopy(raw_hist), None, ""
-    for chunk in stream_generate(raw_hist[:-1], sys_prompt):
+    for chunk in glm4v.stream_generate(raw_hist[:-1], sys_prompt):
         if stop_generation:
             break
         place["content"] = chunk
@@ -247,14 +250,27 @@ def reset():
 
 
 demo = gr.Blocks(title="GLM-4.1V-9B-Thinking", theme=gr.themes.Soft())
+
 with demo:
     gr.Markdown(
-        "<div style='text-align:center;font-size:32px;font-weight:bold;margin-bottom:20px;'>GLM-4.1V-9B-Thinking</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/GLM-4.1V-Thinking'>Github</a></div>")
+        "<div style='text-align:center;font-size:32px;font-weight:bold;margin-bottom:20px;'>GLM-4.1V-9B-Thinking</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/GLM-4.1V-Thinking'>Github</a> |</div>"
+        "<a href='https://arxiv.org/abs/2507.01006'>Paper</a> |</div>"
+        "<a href='https://www.bigmodel.cn/dev/api/visual-reasoning-model/GLM-4.1V-Thinking'>API</a> |</div>"
+    )
     raw_history = gr.State([])
+
     with gr.Row():
         with gr.Column(scale=7):
-            chatbox = gr.Chatbot(label="Chat", type="messages", height=600, elem_classes="chatbot-container",
-                                 sanitize_html=False, line_breaks=True)
+            chatbox = gr.Chatbot(
+                label="Chat",
+                type="messages",
+                height=600,
+                elem_classes="chatbot-container",
+                sanitize_html=False,
+                line_breaks=True
+            )
             textbox = gr.Textbox(label="Message", lines=3)
             with gr.Row():
                 send = gr.Button("Send", variant="primary")
@@ -263,10 +279,25 @@ with demo:
             up = gr.File(label="Upload Files", file_count="multiple", file_types=["file"], type="filepath")
             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.")
+                "The maximum supported input is 10 images or 1 video/PPT/PDF(less than 10 pages) in this demo. "
+                "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, 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])
+
+    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__":
-    demo.launch()
+    demo.launch()