Spaces:
Running
on
Zero
Running
on
Zero
ZhengPeng7
commited on
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.'
|