Add weights option to BiRefNet trained in all different settings.

app.py

@@ -36,11 +36,19 @@ class ImagePreprocessor():
 
 
 from transformers import AutoModelForImageSegmentation
-weights_path = 'zhengpeng7/BiRefNet'
-birefnet = AutoModelForImageSegmentation.from_pretrained(weights_path, trust_remote_code=True)
+weights_path = 'BiRefNet'
+birefnet = AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7', weights_path)), trust_remote_code=True)
 birefnet.to(device)
 birefnet.eval()
 
+usage_to_weights_file = {
+    'General': 'BiRefNet',
+    'Portrait': 'BiRefNet-portrait',
+    'DIS': 'BiRefNet-DIS5K',
+    'HRSOD': 'BiRefNet-HRSOD',
+    'COD': 'BiRefNet-COD',
+    'DIS-TR_TEs': 'BiRefNet-DIS5K-TR_TEs'
+}
 
 # def predict(image_1, image_2):
 #     images = [image_1, image_2]
@@ -50,10 +58,10 @@ def predict(image, resolution, weights_file):
     global birefnet
     if weights_file != weights_path:
         # Load BiRefNet with chosen weights
-        birefnet = AutoModelForImageSegmentation.from_pretrained(weights_file if weights_file is not None else 'zhengpeng7/BiRefNet', trust_remote_code=True)
+        birefnet = AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7', usage_to_weights_file[weights_file] if weights_file is not None else 'BiRefNet')), trust_remote_code=True)
         birefnet.to(device)
         birefnet.eval()
-        # weights_path = weights_file
+        weights_path = weights_file
 
     resolution = f"{image.shape[1]}x{image.shape[0]}" if resolution == '' else resolution
     # Image is a RGB numpy array.
@@ -97,7 +105,7 @@ demo = gr.Interface(
     inputs=[
         'image',
         gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `512x512`. Higher resolutions can be much slower for inference.", label="Resolution"),
-        gr.Radio(['zhengpeng7/BiRefNet', 'zhengpeng7/BiRefNet-portrait'], label="Weights", info="Choose the weights you want.")
+        gr.Radio(list(usage_to_weights_file.keys()), label="Weights", info="Choose the weights you want.")
     ],
     outputs=ImageSlider(),
     examples=examples,