Add the weights option.

app.py

@@ -5,6 +5,7 @@ import numpy as np
 from PIL import Image
 import torch
 from torchvision import transforms
+from transformers import AutoModelForImageSegmentation
 import gradio as gr
 import spaces
 from gradio_imageslider import ImageSlider
@@ -34,9 +35,9 @@ class ImagePreprocessor():
         return image
 
 
-
 from transformers import AutoModelForImageSegmentation
-birefnet = AutoModelForImageSegmentation.from_pretrained('zhengpeng7/BiRefNet', trust_remote_code=True)
+model_path = 'zhengpeng7/BiRefNet'
+birefnet = AutoModelForImageSegmentation.from_pretrained(model_path, trust_remote_code=True)
 birefnet.to(device)
 birefnet.eval()
 
@@ -44,7 +45,12 @@ birefnet.eval()
 # def predict(image_1, image_2):
 #     images = [image_1, image_2]
 @spaces.GPU
-def predict(image, resolution):
+def predict(image, resolution, weights_file):
+    # Load BiRefNet with chosen weights
+    birefnet = AutoModelForImageSegmentation.from_pretrained(weights_file, trust_remote_code=True)
+    birefnet.to(device)
+    birefnet.eval()
+
     resolution = f"{image.shape[1]}x{image.shape[0]}" if resolution == '' else resolution
     # Image is a RGB numpy array.
     resolution = [int(int(reso)//32*32) for reso in resolution.strip().split('x')]
@@ -84,7 +90,11 @@ examples[-1][1] = '512x512'
 
 demo = gr.Interface(
     fn=predict,
-    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")],
+    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.Checkbox(['zhengpeng7/BiRefNet', 'zhengpeng7/BiRefNet-portrait'], label="Models", info="Choose the weights you want.")
+    ],
     outputs=ImageSlider(),
     examples=examples,
     title='Online demo for `Bilateral Reference for High-Resolution Dichotomous Image Segmentation`',