OriLib commited on
Commit
324bfa5
1 Parent(s): fad4871

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +3 -1
README.md CHANGED
@@ -107,13 +107,15 @@ import torch, os
107
  from PIL import Image
108
  from briarmbg import BriaRMBG
109
  from utilities import preprocess_image, postprocess_image
 
110
 
111
- model_path = f"{os.path.dirname(os.path.abspath(__file__))}/model.pth"
112
  im_path = f"{os.path.dirname(os.path.abspath(__file__))}/example_input.jpg"
113
 
114
  net = BriaRMBG()
115
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
116
  net.load_state_dict(torch.load(model_path, map_location=device))
 
117
  net.eval()
118
 
119
  # prepare input
 
107
  from PIL import Image
108
  from briarmbg import BriaRMBG
109
  from utilities import preprocess_image, postprocess_image
110
+ from huggingface_hub import hf_hub_download
111
 
112
+ model_path = hf_hub_download("briaai/RMBG-1.4", 'model.pth')
113
  im_path = f"{os.path.dirname(os.path.abspath(__file__))}/example_input.jpg"
114
 
115
  net = BriaRMBG()
116
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
117
  net.load_state_dict(torch.load(model_path, map_location=device))
118
+ net.to(device)
119
  net.eval()
120
 
121
  # prepare input