ZhengPeng7 commited on
Commit
aa6bac1
·
1 Parent(s): d719afe

Add a new model for matting: BiRefNet_dynamic-matting.

Browse files
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -150,13 +150,14 @@ usage_to_weights_file = {
150
  'General-reso_512': 'BiRefNet_512x512',
151
  'General-Lite': 'BiRefNet_lite',
152
  'General-Lite-2K': 'BiRefNet_lite-2K',
153
- 'Anime-Lite': 'BiRefNet_lite-Anime',
154
  'DIS': 'BiRefNet-DIS5K',
155
  'HRSOD': 'BiRefNet-HRSOD',
156
  'COD': 'BiRefNet-COD',
157
  'DIS-TR_TEs': 'BiRefNet-DIS5K-TR_TEs',
158
  'General-legacy': 'BiRefNet-legacy',
159
  'General-dynamic': 'BiRefNet_dynamic',
 
160
  }
161
 
162
  birefnet = transformers.AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7', usage_to_weights_file['General'])), trust_remote_code=True)
@@ -186,7 +187,7 @@ def predict(images, resolution, weights_file):
186
  elif weights_file in ['General-reso_512']:
187
  resolution = (512, 512)
188
  else:
189
- if weights_file in ['General-dynamic']:
190
  resolution = None
191
  print('Using the original size (div by 32) for inference.')
192
  else:
 
150
  'General-reso_512': 'BiRefNet_512x512',
151
  'General-Lite': 'BiRefNet_lite',
152
  'General-Lite-2K': 'BiRefNet_lite-2K',
153
+ # 'Anime-Lite': 'BiRefNet_lite-Anime',
154
  'DIS': 'BiRefNet-DIS5K',
155
  'HRSOD': 'BiRefNet-HRSOD',
156
  'COD': 'BiRefNet-COD',
157
  'DIS-TR_TEs': 'BiRefNet-DIS5K-TR_TEs',
158
  'General-legacy': 'BiRefNet-legacy',
159
  'General-dynamic': 'BiRefNet_dynamic',
160
+ 'Matting-dynamic': 'BiRefNet_dynamic-matting',
161
  }
162
 
163
  birefnet = transformers.AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7', usage_to_weights_file['General'])), trust_remote_code=True)
 
187
  elif weights_file in ['General-reso_512']:
188
  resolution = (512, 512)
189
  else:
190
+ if '_dynamic' in weights_file:
191
  resolution = None
192
  print('Using the original size (div by 32) for inference.')
193
  else: