Update the model switch part.

app.py

@@ -36,8 +36,8 @@ class ImagePreprocessor():
 
 
 from transformers import AutoModelForImageSegmentation
-model_path = 'zhengpeng7/BiRefNet'
-birefnet = AutoModelForImageSegmentation.from_pretrained(model_path, trust_remote_code=True)
+weights_path = 'zhengpeng7/BiRefNet'
+birefnet = AutoModelForImageSegmentation.from_pretrained(weights_path, trust_remote_code=True)
 birefnet.to(device)
 birefnet.eval()
 
@@ -46,10 +46,13 @@ birefnet.eval()
 #     images = [image_1, image_2]
 @spaces.GPU
 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()
+    global weights_path
+    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.to(device)
+        birefnet.eval()
+        # weights_path = weights_file
 
     resolution = f"{image.shape[1]}x{image.shape[0]}" if resolution == '' else resolution
     # Image is a RGB numpy array.