Update app.py

app.py

@@ -84,28 +84,28 @@ usage_to_weights_file = {
     'General-legacy': 'BiRefNet-legacy'
 }
 
-birefnet = AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7', usage_to_weights_file['General'])), trust_remote_code=True)
+birefnet = AutoModelForImageSegmentation.from_pretrained('/'.join('zhengpeng7', 'BiRefNet_lite'), trust_remote_code=True)
 birefnet.to(device)
 birefnet.eval()
 
 
 @spaces.GPU
-def predict(images, resolution, weights_file):
+def predict(images):
     assert (images is not None), 'AssertionError: images cannot be None.'
 
     global birefnet
     # 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']))
+    _weights_file = '/'.join('zhengpeng7', 'BiRefNet_lite')
     print('Using weights: {}.'.format(_weights_file))
     birefnet = AutoModelForImageSegmentation.from_pretrained(_weights_file, trust_remote_code=True)
     birefnet.to(device)
     birefnet.eval()
 
-    try:
-        resolution = [int(int(reso)//32*32) for reso in resolution.strip().split('x')]
-    except:
-        resolution = (1024, 1024) if weights_file not in ['General-Lite-2K'] else (2560, 1440)
-        print('Invalid resolution input. Automatically changed to 1024x1024 or 2K.')
+    #try:
+    #    resolution = [int(int(reso)//32*32) for reso in resolution.strip().split('x')]
+    #except:
+    #    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
@@ -131,7 +131,7 @@ def predict(images, resolution, weights_file):
 
         image = image_ori.convert('RGB')
         # Preprocess the image
-        image_preprocessor = ImagePreprocessor(resolution=tuple(resolution))
+        image_preprocessor = ImagePreprocessor() #(resolution=tuple(resolution))
         image_proc = image_preprocessor.proc(image)
         image_proc = image_proc.unsqueeze(0)
 
@@ -184,44 +184,19 @@ 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`.", label="Resolution"),
-        gr.Radio(list(usage_to_weights_file.keys()), value='General', label="Weights", info="Choose the weights you want.")
+        #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"),
-    examples=examples,
+    #examples=examples,
     api_name="image",
     description=descriptions,
 )
 
-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`.", 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"),
-    examples=examples_url,
-    api_name="text",
-    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(
-    [tab_image, tab_text, tab_batch],
-    ['image', 'text', 'batch'],
-    title="BiRefNet demo for subject extraction (general / matting / salient / camouflaged / portrait).",
+    [tab_image],
+    ['image'],
+    title="BiRefNet demo for subject extraction.",
 )
 
 if __name__ == "__main__":