kvinod15 commited on
Commit
1e42f2e
·
verified ·
1 Parent(s): e86315f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +9 -1
app.py CHANGED
@@ -1,3 +1,4 @@
 
1
  import io
2
  import numpy as np
3
  import torch
@@ -13,10 +14,17 @@ from transformers import AutoModelForImageSegmentation, pipeline
13
  # Set device (GPU if available, else CPU)
14
  device = "cuda" if torch.cuda.is_available() else "cpu"
15
 
 
 
 
 
 
16
  # Load the segmentation model (RMBG-2.0)
 
17
  segmentation_model = AutoModelForImageSegmentation.from_pretrained(
18
  'briaai/RMBG-2.0',
19
- trust_remote_code=True
 
20
  )
21
  segmentation_model.to(device)
22
  segmentation_model.eval()
 
1
+ import os
2
  import io
3
  import numpy as np
4
  import torch
 
14
  # Set device (GPU if available, else CPU)
15
  device = "cuda" if torch.cuda.is_available() else "cpu"
16
 
17
+ # Retrieve Hugging Face token from environment variable
18
+ hf_token = os.environ.get("HF_ACCESS_TOKEN")
19
+ if hf_token is None:
20
+ print("Warning: HF_ACCESS_TOKEN is not set. Access to gated models may fail.")
21
+
22
  # Load the segmentation model (RMBG-2.0)
23
+ # Ensure you have access to this gated repo and that your token is valid.
24
  segmentation_model = AutoModelForImageSegmentation.from_pretrained(
25
  'briaai/RMBG-2.0',
26
+ trust_remote_code=True,
27
+ use_auth_token=hf_token
28
  )
29
  segmentation_model.to(device)
30
  segmentation_model.eval()