Update and sync the codes for running locally.

app_local.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
@@ -60,8 +60,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]),
         ])
@@ -77,17 +78,19 @@ usage_to_weights_file = {
     'Matting-HR': 'BiRefNet_HR-matting',
     'Matting': 'BiRefNet-matting',
     'Portrait': 'BiRefNet-portrait',
-    'General-reso_512': 'BiRefNet-reso_512',
+    'General-reso_512': 'BiRefNet_512x512',
     'General-Lite': 'BiRefNet_lite',
     'General-Lite-2K': 'BiRefNet_lite-2K',
+    'Anime-Lite': 'BiRefNet_lite-Anime',
     'DIS': 'BiRefNet-DIS5K',
     '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)
@@ -169,7 +180,7 @@ def predict(images, resolution, weights_file):
                 zipf.write(file, os.path.basename(file))
         return save_paths, zip_file_path
     else:
-        return (image_masked, image_ori)[0]
+        return (image_masked, image_ori)
 
 
 examples = [[_] for _ in glob('examples/*')][:]
@@ -201,7 +212,7 @@ tab_image = gr.Interface(
         gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`.", label="Resolution"),
         gr.Radio(list(usage_to_weights_file.keys()), value='General', label="Weights", info="Choose the weights you want.")
     ],
-    outputs=gr.Image(label="BiRefNet's prediction", type="pil", format='png'),
+    outputs=gr.ImageSlider(label="BiRefNet's prediction", type="pil", format='png'),
     examples=examples,
     api_name="image",
     description=descriptions,
@@ -214,7 +225,7 @@ tab_text = gr.Interface(
         gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`.", label="Resolution"),
         gr.Radio(list(usage_to_weights_file.keys()), value='General', label="Weights", info="Choose the weights you want.")
     ],
-    outputs=gr.Image(label="BiRefNet's prediction", type="pil", format='png'),
+    outputs=gr.ImageSlider(label="BiRefNet's prediction", type="pil", format='png'),
     examples=examples_url,
     api_name="URL",
     description=descriptions+'\nTab-URL is partially modified from https://huggingface.co/spaces/not-lain/background-removal, thanks to this great work!',
@@ -235,7 +246,7 @@ tab_batch = gr.Interface(
 demo = gr.TabbedInterface(
     [tab_image, tab_text, tab_batch],
     ['image', 'URL', 'batch'],
-    title="BiRefNet demo for subject extraction and background removal ([CAAI AIR'24] Bilateral Reference for High-Resolution Dichotomous Image Segmentation).",
+    title="Official Online Demo of BiRefNet",
 )
 
 if __name__ == "__main__":