nihuajian commited on
Commit
914233b
·
verified ·
1 Parent(s): 8ecc3b6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +9 -10
app.py CHANGED
@@ -4,6 +4,7 @@ from PIL import Image
4
  from transformers import AutoModel, AutoTokenizer
5
  import warnings
6
  import os
 
7
 
8
  # 禁用警告信息
9
  warnings.filterwarnings("ignore")
@@ -12,21 +13,22 @@ warnings.filterwarnings("ignore")
12
  model = None
13
  tokenizer = None
14
 
 
15
  def load_model():
16
  """加载MiniCPM-o模型"""
17
  global model, tokenizer
18
  if model is None:
19
  print("正在加载MiniCPM-o模型...")
 
20
  model = AutoModel.from_pretrained(
21
  'openbmb/MiniCPM-o-2_6',
22
  trust_remote_code=True,
23
- attn_implementation='sdpa',
24
- torch_dtype=torch.bfloat16,
25
  init_vision=True,
26
  init_audio=False,
27
  init_tts=False
28
  )
29
- model = model.eval().cuda()
30
  tokenizer = AutoTokenizer.from_pretrained('openbmb/MiniCPM-o-2_6', trust_remote_code=True)
31
  print("模型加载完成")
32
  return model, tokenizer
@@ -133,6 +135,7 @@ def clean_text_output(text):
133
 
134
  return '\n'.join(text_lines)
135
 
 
136
  def parse_image(image, parse_type):
137
  """解析图片内容为指定格式"""
138
  try:
@@ -298,13 +301,9 @@ def create_interface():
298
  return interface
299
 
300
  if __name__ == "__main__":
301
- # 预加载模型(可选,在启动时加载以减少首次使用延迟)
302
- try:
303
- load_model()
304
- print("✅ 模型预加载完成")
305
- except Exception as e:
306
- print(f"⚠️ 模型预加载失败: {e}")
307
- print("模型将在首次使用时加载")
308
 
309
  # 创建并启动界面
310
  interface = create_interface()
 
4
  from transformers import AutoModel, AutoTokenizer
5
  import warnings
6
  import os
7
+ import spaces
8
 
9
  # 禁用警告信息
10
  warnings.filterwarnings("ignore")
 
13
  model = None
14
  tokenizer = None
15
 
16
+ @spaces.GPU
17
  def load_model():
18
  """加载MiniCPM-o模型"""
19
  global model, tokenizer
20
  if model is None:
21
  print("正在加载MiniCPM-o模型...")
22
+ device = "cuda" if torch.cuda.is_available() else "cpu"
23
  model = AutoModel.from_pretrained(
24
  'openbmb/MiniCPM-o-2_6',
25
  trust_remote_code=True,
26
+ torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
 
27
  init_vision=True,
28
  init_audio=False,
29
  init_tts=False
30
  )
31
+ model = model.eval().to(device)
32
  tokenizer = AutoTokenizer.from_pretrained('openbmb/MiniCPM-o-2_6', trust_remote_code=True)
33
  print("模型加载完成")
34
  return model, tokenizer
 
135
 
136
  return '\n'.join(text_lines)
137
 
138
+ @spaces.GPU
139
  def parse_image(image, parse_type):
140
  """解析图片内容为指定格式"""
141
  try:
 
301
  return interface
302
 
303
  if __name__ == "__main__":
304
+ # 在ZeroGPU环境中不预加载模型,按需加载以节省资源
305
+ print("🚀 启动MiniCPM多模态内容解析工具")
306
+ print("📝 模型将在首次使用时自动加载")
 
 
 
 
307
 
308
  # 创建并启动界面
309
  interface = create_interface()