ZhengPeng7 commited on
Commit
21603dc
1 Parent(s): 0e4e344

Update the weights to the one for general use with massive training.

Browse files
Files changed (1) hide show
  1. app.py +1 -1
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.'