Add BiRefNet-Lite-2K model, add model choosing option in tab_batch.

app.py

@@ -61,7 +61,7 @@ def FB_blur_fusion_foreground_estimator(image, F, B, alpha, r=90):
 class ImagePreprocessor():
     def __init__(self, resolution: Tuple[int, int] = (1024, 1024)) -> None:
         self.transform_image = transforms.Compose([
-            transforms.Resize(resolution),    # 1. keep consistent with the cv2.resize used in training 2. redundant with that in path_to_image()
+            transforms.Resize(resolution),
             transforms.ToTensor(),
             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
         ])
@@ -74,6 +74,7 @@ class ImagePreprocessor():
 usage_to_weights_file = {
     'General': 'BiRefNet',
     'General-Lite': 'BiRefNet_lite',
+    'General-Lite-2K': 'BiRefNet_lite-2K',
     'Portrait': 'BiRefNet-portrait',
     'DIS': 'BiRefNet-DIS5K',
     'HRSOD': 'BiRefNet-HRSOD',
@@ -102,8 +103,8 @@ def predict(images, resolution, weights_file):
     try:
         resolution = [int(int(reso)//32*32) for reso in resolution.strip().split('x')]
     except:
-        resolution = [1024, 1024]
-        print('Invalid resolution input. Automatically changed to 1024x1024.')
+        resolution = (1024, 1024) if weights_file not in ['General-Lite-2K'] else (2560, 1440)
+        print('Invalid resolution input. Automatically changed to 1024x1024 or 2K.')
 
     if isinstance(images, list):
         # For tab_batch
@@ -182,7 +183,7 @@ tab_image = gr.Interface(
     fn=predict,
     inputs=[
         gr.Image(label='Upload an image'),
-        gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`. Higher resolutions can be much slower for inference.", label="Resolution"),
+        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=ImageSlider(label="BiRefNet's prediction", type="pil"),
@@ -195,7 +196,7 @@ tab_text = gr.Interface(
     fn=predict,
     inputs=[
         gr.Textbox(label="Paste an image URL"),
-        gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`. Higher resolutions can be much slower for inference.", label="Resolution"),
+        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=ImageSlider(label="BiRefNet's prediction", type="pil"),
@@ -206,7 +207,11 @@ tab_text = gr.Interface(
 
 tab_batch = gr.Interface(
     fn=predict,
-    inputs=gr.File(label="Upload multiple images", type="filepath", file_count="multiple"),
+    inputs=[
+        gr.File(label="Upload multiple images", type="filepath", file_count="multiple"),
+        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.Gallery(label="BiRefNet's predictions"), gr.File(label="Download masked images.")],
     api_name="batch",
     description=descriptions+'\nTab-batch is partially modified from https://huggingface.co/spaces/NegiTurkey/Multi_Birefnetfor_Background_Removal, thanks to this great work!',

app_local.py

@@ -61,7 +61,7 @@ def FB_blur_fusion_foreground_estimator(image, F, B, alpha, r=90):
 class ImagePreprocessor():
     def __init__(self, resolution: Tuple[int, int] = (1024, 1024)) -> None:
         self.transform_image = transforms.Compose([
-            transforms.Resize(resolution),    # 1. keep consistent with the cv2.resize used in training 2. redundant with that in path_to_image()
+            transforms.Resize(resolution),
             transforms.ToTensor(),
             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
         ])
@@ -74,6 +74,7 @@ class ImagePreprocessor():
 usage_to_weights_file = {
     'General': 'BiRefNet',
     'General-Lite': 'BiRefNet_lite',
+    'General-Lite-2K': 'BiRefNet_lite-2K',
     'Portrait': 'BiRefNet-portrait',
     'DIS': 'BiRefNet-DIS5K',
     'HRSOD': 'BiRefNet-HRSOD',
@@ -102,8 +103,8 @@ def predict(images, resolution, weights_file):
     try:
         resolution = [int(int(reso)//32*32) for reso in resolution.strip().split('x')]
     except:
-        resolution = [1024, 1024]
-        print('Invalid resolution input. Automatically changed to 1024x1024.')
+        resolution = (1024, 1024) if weights_file not in ['General-Lite-2K'] else (2560, 1440)
+        print('Invalid resolution input. Automatically changed to 1024x1024 or 2K.')
 
     if isinstance(images, list):
         # For tab_batch
@@ -182,7 +183,7 @@ tab_image = gr.Interface(
     fn=predict,
     inputs=[
         gr.Image(label='Upload an image'),
-        gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`. Higher resolutions can be much slower for inference.", label="Resolution"),
+        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'),
@@ -195,7 +196,7 @@ tab_text = gr.Interface(
     fn=predict,
     inputs=[
         gr.Textbox(label="Paste an image URL"),
-        gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`. Higher resolutions can be much slower for inference.", label="Resolution"),
+        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'),
@@ -206,7 +207,11 @@ tab_text = gr.Interface(
 
 tab_batch = gr.Interface(
     fn=predict,
-    inputs=gr.File(label="Upload multiple images", type="filepath", file_count="multiple"),
+    inputs=[
+        gr.File(label="Upload multiple images", type="filepath", file_count="multiple"),
+        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.Gallery(label="BiRefNet's predictions"), gr.File(label="Download masked images.")],
     api_name="batch",
     description=descriptions+'\nTab-batch is partially modified from https://huggingface.co/spaces/NegiTurkey/Multi_Birefnetfor_Background_Removal, thanks to this great work!',