Update README.md

README.md

@@ -107,13 +107,15 @@ import torch, os
 from PIL import Image
 from briarmbg import BriaRMBG
 from utilities import preprocess_image, postprocess_image
+from huggingface_hub import hf_hub_download
 
-model_path = f"{os.path.dirname(os.path.abspath(__file__))}/model.pth"
+model_path = hf_hub_download("briaai/RMBG-1.4", 'model.pth')
 im_path = f"{os.path.dirname(os.path.abspath(__file__))}/example_input.jpg"
 
 net = BriaRMBG()
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 net.load_state_dict(torch.load(model_path, map_location=device))
+net.to(device)
 net.eval()    
 
 # prepare input