Update code snippet
Browse files`from_pretrained` puts the model in evaluation mode by default, and `hf_hub_download` is no longer needed.
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]
|