XiaoyiYangRIT commited on
Commit
62b60d1
·
1 Parent(s): 741cc94

Update some files

Browse files
Files changed (1) hide show
  1. app.py +13 -4
app.py CHANGED
@@ -1,6 +1,7 @@
1
  import gradio as gr
2
  import torch
3
  import math
 
4
  from transformers import AutoTokenizer, AutoModel, AutoProcessor
5
  from decord import VideoReader, cpu
6
  from PIL import Image
@@ -17,10 +18,18 @@ transform = Compose([
17
  ])
18
 
19
  # === 模型加载 ===
 
20
  MODEL_NAME = "OpenGVLab/InternVL3-14B"
21
 
22
- tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
23
- processor = AutoProcessor.from_pretrained(MODEL_NAME, trust_remote_code=True)
 
 
 
 
 
 
 
24
 
25
  def split_model(model_path):
26
  from transformers import AutoConfig
@@ -47,10 +56,10 @@ def split_model(model_path):
47
  device_map[f'language_model.model.layers.{num_layers - 1}'] = 0
48
  return device_map
49
 
50
- device_map = split_model(MODEL_NAME)
51
 
52
  model = AutoModel.from_pretrained(
53
- MODEL_NAME,
54
  torch_dtype=torch.bfloat16,
55
  low_cpu_mem_usage=True,
56
  use_flash_attn=True,
 
1
  import gradio as gr
2
  import torch
3
  import math
4
+ import os
5
  from transformers import AutoTokenizer, AutoModel, AutoProcessor
6
  from decord import VideoReader, cpu
7
  from PIL import Image
 
18
  ])
19
 
20
  # === 模型加载 ===
21
+ PERSISTENT_DIR = "/data/internvl3_model" # 持久路径
22
  MODEL_NAME = "OpenGVLab/InternVL3-14B"
23
 
24
+ # 如果第一次运行:下载模型并缓存到 /data
25
+ if not os.path.exists(PERSISTENT_DIR):
26
+ print("Downloading model to persistent storage...")
27
+ from transformers import snapshot_download
28
+ snapshot_download(repo_id=MODEL_NAME, local_dir=PERSISTENT_DIR, trust_remote_code=True)
29
+
30
+ # 模型加载(从本地)
31
+ tokenizer = AutoTokenizer.from_pretrained(PERSISTENT_DIR, trust_remote_code=True)
32
+ processor = AutoProcessor.from_pretrained(PERSISTENT_DIR, trust_remote_code=True)
33
 
34
  def split_model(model_path):
35
  from transformers import AutoConfig
 
56
  device_map[f'language_model.model.layers.{num_layers - 1}'] = 0
57
  return device_map
58
 
59
+ device_map = split_model(PERSISTENT_DIR)
60
 
61
  model = AutoModel.from_pretrained(
62
+ PERSISTENT_DIR,
63
  torch_dtype=torch.bfloat16,
64
  low_cpu_mem_usage=True,
65
  use_flash_attn=True,