Update app.py

app.py

@@ -5,7 +5,6 @@ import torch
 import gradio as gr
 import spaces
 
-from glob import glob
 from typing import Tuple
 
 from PIL import Image
@@ -15,7 +14,7 @@ import requests
 from io import BytesIO
 import zipfile
 
-# Fix the HF space permission error when using from_pretrained(..., trust_remote_code=True)
+# Fix the HF space permission error
 os.environ["HF_MODULES_CACHE"] = os.path.join("/tmp/hf_cache", "modules")
 
 import transformers
@@ -26,7 +25,6 @@ torch.jit.script = lambda f: f
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-### image_proc.py
 def refine_foreground(image, mask, r=90):
     if mask.size != image.size:
         mask = mask.resize(image.size)
@@ -66,58 +64,26 @@ class ImagePreprocessor():
     def proc(self, image: Image.Image) -> torch.Tensor:
         return self.transform_image(image)
 
-usage_to_weights_file = {
-    'General': 'BiRefNet',
-    'General-HR': 'BiRefNet_HR',
-    'Matting-HR': 'BiRefNet_HR-matting',
-    'Matting': 'BiRefNet-matting',
-    'Portrait': 'BiRefNet-portrait',
-    '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-dynamic': 'BiRefNet_dynamic',
-}
-
+# Fixed weights
+weights_file = 'BiRefNet'
 birefnet = transformers.AutoModelForImageSegmentation.from_pretrained(
-    '/'.join(('zhengpeng7', usage_to_weights_file['General'])),
-    trust_remote_code=True
+    '/'.join(('zhengpeng7', weights_file)), trust_remote_code=True
 )
 birefnet.to(device)
 birefnet.eval(); birefnet.half()
 
 @spaces.GPU
-def predict(images, resolution, weights_file):
+def predict(images, resolution):
     assert images is not None, 'AssertionError: images cannot be None.'
 
-    global birefnet
-    _weights_file = '/'.join(('zhengpeng7', usage_to_weights_file[weights_file] if weights_file is not None else usage_to_weights_file['General']))
+    _weights_file = '/'.join(('zhengpeng7', weights_file))
     print('Using weights: {}.'.format(_weights_file))
-    birefnet = transformers.AutoModelForImageSegmentation.from_pretrained(_weights_file, trust_remote_code=True)
-    birefnet.to(device)
-    birefnet.eval(); birefnet.half()
 
     try:
         resolution = [int(int(reso)//32*32) for reso in resolution.strip().split('x')]
     except:
-        if weights_file in ['General-HR', 'Matting-HR']:
-            resolution = (2048, 2048)
-        elif weights_file in ['General-Lite-2K']:
-            resolution = (2560, 1440)
-        elif weights_file in ['General-reso_512']:
-            resolution = (512, 512)
-        else:
-            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.')
+        resolution = (1024, 1024)
+        print('Invalid resolution input. Automatically changed to 1024x1024.')
 
     if isinstance(images, list):
         save_paths = []
@@ -143,8 +109,7 @@ def predict(images, resolution, weights_file):
         image = image_ori.convert('RGB')
         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
+            resolution = resolution_div_by_32
         image_preprocessor = ImagePreprocessor(resolution=tuple(resolution))
         image_proc = image_preprocessor.proc(image).unsqueeze(0)
 
@@ -182,7 +147,6 @@ tab_image = gr.Interface(
     inputs=[
         gr.Image(label='Upload an image'),
         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.ImageSlider(label="BiRefNet's prediction", type="pil", format='png'),
     api_name="image",
@@ -194,21 +158,21 @@ tab_text = gr.Interface(
     inputs=[
         gr.Textbox(label="Paste an image URL"),
         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.ImageSlider(label="BiRefNet's prediction", type="pil", format='png'),
     api_name="URL",
-    )
+    description=descriptions + '\nTab-URL is partially modified from https://huggingface.co/spaces/not-lain/background-removal, thanks to this great work!',
+)
 
 tab_batch = gr.Interface(
     fn=predict,
     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!',
 )
 
 demo = gr.TabbedInterface(