Spaces:
Running
on
Zero
Running
on
Zero
Commit
•
21603dc
1
Parent(s):
0e4e344
Update the weights to the one for general use with massive training.
Browse files
app.py
CHANGED
|
@@ -40,7 +40,7 @@ class ImagePreprocessor():
|
|
| 40 |
|
| 41 |
|
| 42 |
model = BiRefNet(bb_pretrained=False)
|
| 43 |
-
state_dict = './BiRefNet_ep580.pth'
|
| 44 |
if os.path.exists(state_dict):
|
| 45 |
birefnet_dict = torch.load(state_dict, map_location="cpu")
|
| 46 |
unwanted_prefix = '_orig_mod.'
|
|
|
|
| 40 |
|
| 41 |
|
| 42 |
model = BiRefNet(bb_pretrained=False)
|
| 43 |
+
state_dict = ['./BiRefNet_ep580.pth', 'BiRefNet-massive-epoch_240.pth'][-1]
|
| 44 |
if os.path.exists(state_dict):
|
| 45 |
birefnet_dict = torch.load(state_dict, map_location="cpu")
|
| 46 |
unwanted_prefix = '_orig_mod.'
|