Fly-ShuAI commited on
Commit
12c3012
·
verified ·
1 Parent(s): 948456e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -0
app.py CHANGED
@@ -15,6 +15,14 @@ num_frames, width, height = 49, 832, 480
15
  gpu_id = 0
16
  device = f'cuda:{gpu_id}'
17
 
 
 
 
 
 
 
 
 
18
  rmbg_model = AutoModelForImageSegmentation.from_pretrained('briaai/RMBG-2.0', trust_remote_code=True)
19
  torch.set_float32_matmul_precision(['high', 'highest'][0])
20
  rmbg_model.to(device)
 
15
  gpu_id = 0
16
  device = f'cuda:{gpu_id}'
17
 
18
+ from huggingface_hub import login
19
+ import os
20
+ # 使用环境变量中的 token 登录
21
+ token = os.environ.get("HF_TOKEN")
22
+ if token:
23
+ login(token)
24
+
25
+
26
  rmbg_model = AutoModelForImageSegmentation.from_pretrained('briaai/RMBG-2.0', trust_remote_code=True)
27
  torch.set_float32_matmul_precision(['high', 'highest'][0])
28
  rmbg_model.to(device)