Turn back back to the latest version.

app.py

@@ -10,7 +10,7 @@ from typing import Tuple
 
 from PIL import Image
 from gradio_imageslider import ImageSlider
-from transformers import AutoModelForImageSegmentation
+import transformers
 from torchvision import transforms
 
 import requests
@@ -18,6 +18,7 @@ from io import BytesIO
 import zipfile
 
 
+transformers.utils.move_cache()
 torch.set_float32_matmul_precision('high')
 torch.jit.script = lambda f: f
 
@@ -60,8 +61,9 @@ def FB_blur_fusion_foreground_estimator(image, F, B, alpha, r=90):
 
 class ImagePreprocessor():
     def __init__(self, resolution: Tuple[int, int] = (1024, 1024)) -> None:
+        # Input resolution is on WxH.
         self.transform_image = transforms.Compose([
-            transforms.Resize(resolution),
+            transforms.Resize(resolution[::-1]),
             transforms.ToTensor(),
             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
         ])
@@ -84,10 +86,11 @@ usage_to_weights_file = {
     'HRSOD': 'BiRefNet-HRSOD',
     'COD': 'BiRefNet-COD',
     'DIS-TR_TEs': 'BiRefNet-DIS5K-TR_TEs',
-    'General-legacy': 'BiRefNet-legacy'
+    'General-legacy': 'BiRefNet-legacy',
+    'General-dynamic': 'BiRefNet_dynamic',
 }
 
-birefnet = AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7', usage_to_weights_file['General'])), trust_remote_code=True)
+birefnet = transformers.AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7', usage_to_weights_file['General'])), trust_remote_code=True)
 birefnet.to(device)
 birefnet.eval(); birefnet.half()
 
@@ -100,7 +103,7 @@ def predict(images, resolution, weights_file):
     # Load BiRefNet with chosen weights
     _weights_file = '/'.join(('zhengpeng7', usage_to_weights_file[weights_file] if weights_file is not None else usage_to_weights_file['General']))
     print('Using weights: {}.'.format(_weights_file))
-    birefnet = AutoModelForImageSegmentation.from_pretrained(_weights_file, trust_remote_code=True)
+    birefnet = transformers.AutoModelForImageSegmentation.from_pretrained(_weights_file, trust_remote_code=True)
     birefnet.to(device)
     birefnet.eval(); birefnet.half()
 
@@ -114,7 +117,11 @@ def predict(images, resolution, weights_file):
         elif weights_file in ['General-reso_512']:
             resolution = (512, 512)
         else:
-            resolution = (1024, 1024)
+            if weights_file in ['General-dynamic']:
+                resolution = None
+                print('Using the original size (div by 32) for inference.')
+            else:
+                resolution = (1024, 1024)
         print('Invalid resolution input. Automatically changed to 1024x1024 / 2048x2048 / 2560x1440.')
 
     if isinstance(images, list):
@@ -141,6 +148,10 @@ def predict(images, resolution, weights_file):
 
         image = image_ori.convert('RGB')
         # Preprocess the image
+        if resolution is None:
+            resolution_div_by_32 = [int(int(reso)//32*32) for reso in image.size]
+            if resolution_div_by_32 != resolution:
+                resolution = resolution_div_by_32
         image_preprocessor = ImagePreprocessor(resolution=tuple(resolution))
         image_proc = image_preprocessor.proc(image)
         image_proc = image_proc.unsqueeze(0)