Update README.md
Browse files
README.md
CHANGED
@@ -119,9 +119,9 @@ from PIL import Image
|
|
119 |
import matplotlib.pyplot as plt
|
120 |
import torch
|
121 |
from torchvision import transforms
|
122 |
-
from
|
123 |
|
124 |
-
birefnet =
|
125 |
torch.set_float32_matmul_precision(['high', 'highest'][0])
|
126 |
birefnet.to('cuda')
|
127 |
birefnet.eval()
|
|
|
119 |
import matplotlib.pyplot as plt
|
120 |
import torch
|
121 |
from torchvision import transforms
|
122 |
+
from transformers import AutoModelForImageSegmentation
|
123 |
|
124 |
+
birefnet = AutoModelForImageSegmentation.from_pretrained('briaai/RMBG-2.0', trust_remote_code=True)
|
125 |
torch.set_float32_matmul_precision(['high', 'highest'][0])
|
126 |
birefnet.to('cuda')
|
127 |
birefnet.eval()
|