7

app.py

@@ -15,20 +15,23 @@ import time
 MODEL_PATH = "THUDM/GLM-4.1V-9B-Thinking"
 stop_generation = False
 
+processor = None
+model = None
+
+def load_model():
+    """加载模型和处理器"""
+    global processor, model
+    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",
+    )
 
 class GLM4VModel:
     def __init__(self):
-        self.processor = None
-        self.model = None
-
-    def load(self):
-        self.processor = AutoProcessor.from_pretrained(MODEL_PATH, use_fast=True)
-        self.model = Glm4vForConditionalGeneration.from_pretrained(
-            MODEL_PATH,
-            torch_dtype=torch.bfloat16,
-            device_map="auto",
-            attn_implementation="sdpa",
-        )
+        pass
 
     def _strip_html(self, t):
         return re.sub(r"<[^>]+>", "", t).strip()
@@ -125,19 +128,19 @@ class GLM4VModel:
 
     @spaces.GPU(duration=240)
     def stream_generate(self, raw_hist, sys_prompt):
-        global stop_generation
+        global stop_generation, processor, model
         stop_generation = False
         msgs = self._build_messages(raw_hist, sys_prompt)
-        inputs = self.processor.apply_chat_template(
+        inputs = processor.apply_chat_template(
             msgs,
             tokenize=True,
             add_generation_prompt=True,
             return_dict=True,
             return_tensors="pt",
             padding=True,
-        ).to(self.model.device)
+        ).to(model.device)
 
-        streamer = TextIteratorStreamer(self.processor.tokenizer, skip_prompt=True, skip_special_tokens=False)
+        streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=False)
         gen_args = dict(
             inputs,
             max_new_tokens=8192,
@@ -149,7 +152,7 @@ class GLM4VModel:
             streamer=streamer,
         )
 
-        generation_thread = threading.Thread(target=self.model.generate, kwargs=gen_args)
+        generation_thread = threading.Thread(target=model.generate, kwargs=gen_args)
         generation_thread.start()
 
         buf = ""
@@ -190,8 +193,9 @@ def create_display_history(raw_hist):
     return display_hist
 
 
+# 加载模型和处理器
+load_model()
 glm4v = GLM4VModel()
-glm4v.load()
 
 
 def check_files(files):
@@ -310,4 +314,4 @@ with demo:
     clear.click(reset, outputs=[chatbox, raw_history, up, textbox])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()