nielsr HF staff commited on
Commit
12143b7
1 Parent(s): fa91375

Update code snippet

Browse files

`from_pretrained` puts the model in evaluation mode by default, and `hf_hub_download` is no longer needed.

Files changed (1) hide show
  1. README.md +0 -2
README.md CHANGED
@@ -107,7 +107,6 @@ import torch, os
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
  im_path = f"{os.path.dirname(os.path.abspath(__file__))}/example_input.jpg"
113
 
@@ -115,7 +114,6 @@ net = BriaRMBG()
115
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
116
  net = BriaRMBG.from_pretrained("briaai/RMBG-1.4")
117
  net.to(device)
118
- net.eval()
119
 
120
  # prepare input
121
  model_input_size = [1024,1024]
 
107
  from PIL import Image
108
  from briarmbg import BriaRMBG
109
  from utilities import preprocess_image, postprocess_image
 
110
 
111
  im_path = f"{os.path.dirname(os.path.abspath(__file__))}/example_input.jpg"
112
 
 
114
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
115
  net = BriaRMBG.from_pretrained("briaai/RMBG-1.4")
116
  net.to(device)
 
117
 
118
  # prepare input
119
  model_input_size = [1024,1024]