Update app.py

app.py

@@ -115,13 +115,13 @@ engine = None  # InferEngine 对象
 os.makedirs(LOCAL_MODEL_PATH, exist_ok=True)
 
 def load_llm_model():
-    """使用 ms-swift 的 PtEngine 加载 ScriptAgent 模型"""
+    """使用 ms-swift 的 PtEngine 加载量化模型（int4/int8）"""
     global engine
     if engine is not None:
         return
     
     try:
-        # 1. 检查本地是否已下载模型
+        # 检查模型是否存在
         if not os.path.exists(os.path.join(LOCAL_MODEL_PATH, "config.json")):
             print(f"正在从 HuggingFace 下载模型到 {LOCAL_MODEL_PATH}...")
             snapshot_download(
@@ -134,21 +134,25 @@ def load_llm_model():
         else:
             print(f"✅ 模型已存在: {LOCAL_MODEL_PATH}")
         
-        # 2. 使用 ms-swift 的 PtEngine 加载模型
-        print("正在使用 ms-swift PtEngine 加载模型...")
+        # 🔥 关键修改：使用量化 + 内存优化
+        print("正在使用 ms-swift PtEngine 加载模型（量化模式）...")
         
-        # 🔥 关键修改：使用 PtEngine
         engine = PtEngine(
             model_id_or_path=LOCAL_MODEL_PATH,
-            torch_dtype=torch.float16,      # 半精度
-            max_batch_size=1,                # 批处理大小
-            device_map='cpu',                # CPU设备
+            torch_dtype=torch.float16,
+            max_batch_size=1,
+            device_map='cpu',
+            # 🔥 添加量化参数
+            quant_method='bnb',           # 使用 bitsandbytes 量化
+            quant_bits=4,           # 4-bit 量化（也可以用8）
+            bnb_4bit_compute_dtype=torch.float16,
             model_kwargs={
-                'low_cpu_mem_usage': True,   # 低内存模式
+                'low_cpu_mem_usage': True,
+                'max_memory': {'cpu': '12GB'},  # 限制最大内存使用
             }
         )
         
-        print("✅ ms-swift PtEngine 加载完成")
+        print("✅ ms-swift PtEngine 加载完成（已启用量化）")
         
     except Exception as e:
         print(f"❌ 模型加载失败: {e}")