ZhengPeng7 commited on
Commit
e000df2
·
1 Parent(s): 66fad2f

Add weights option to BiRefNet trained in all different settings.

Browse files
Files changed (1) hide show
  1. app.py +1 -0
app.py CHANGED
@@ -59,6 +59,7 @@ def predict(image, resolution, weights_file):
59
  # Load BiRefNet with chosen weights
60
  _weights_file = '/'.join(('zhengpeng7', usage_to_weights_file[weights_file] if weights_file is not None else 'BiRefNet'))
61
  print('Change weights to:', _weights_file)
 
62
  birefnet = AutoModelForImageSegmentation.from_pretrained(_weights_file, trust_remote_code=True)
63
  birefnet.to(device)
64
  birefnet.eval()
 
59
  # Load BiRefNet with chosen weights
60
  _weights_file = '/'.join(('zhengpeng7', usage_to_weights_file[weights_file] if weights_file is not None else 'BiRefNet'))
61
  print('Change weights to:', _weights_file)
62
+ print('\t', weights_file, weights_path)
63
  birefnet = AutoModelForImageSegmentation.from_pretrained(_weights_file, trust_remote_code=True)
64
  birefnet.to(device)
65
  birefnet.eval()