ZhengPeng7 commited on
Commit
0a4040e
·
verified ·
1 Parent(s): e47ae80

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -41,7 +41,7 @@ class ImagePreprocessor():
41
 
42
 
43
  model = BiRefNet(bb_pretrained=False)
44
- state_dict = ['./BiRefNet_ep590.pth', 'BiRefNet-massive-epoch_240.pth'][0]
45
  if os.path.exists(state_dict):
46
  birefnet_dict = torch.load(state_dict, map_location="cpu")
47
  unwanted_prefix = '_orig_mod.'
 
41
 
42
 
43
  model = BiRefNet(bb_pretrained=False)
44
+ state_dict = ['./BiRefNet-DIS-epoch_590.pth', 'BiRefNet-massive-epoch_240.pth'][0]
45
  if os.path.exists(state_dict):
46
  birefnet_dict = torch.load(state_dict, map_location="cpu")
47
  unwanted_prefix = '_orig_mod.'