ZhengPeng7 commited on
Commit
0e5a7e4
1 Parent(s): bfe6e38

Upgrade the weights loading method to avoid duplicated loading.

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -54,7 +54,7 @@ birefnet.eval()
54
  def predict(image, resolution, weights_file):
55
  global birefnet
56
  # Load BiRefNet with chosen weights
57
- _weights_file = '/'.join(('zhengpeng7', usage_to_weights_file[weights_file] if weights_file is not None else 'BiRefNet'))
58
  print('Change weights to:', _weights_file)
59
  birefnet = birefnet.from_pretrained(_weights_file)
60
  birefnet.to(device)
 
54
  def predict(image, resolution, weights_file):
55
  global birefnet
56
  # Load BiRefNet with chosen weights
57
+ _weights_file = '/'.join(('zhengpeng7', usage_to_weights_file[weights_file] if weights_file is not None else usage_to_weights_file['General']))
58
  print('Change weights to:', _weights_file)
59
  birefnet = birefnet.from_pretrained(_weights_file)
60
  birefnet.to(device)