ZeroGPU and ImageSlider (#1)

- ZeroGPU and ImageSlider (dd31e60100ef70bff8cad77af8ea1f4460a9320d)


Co-authored-by: hysts <[email protected]>

app.py

@@ -6,6 +6,10 @@ from PIL import Image
 import torch
 from torchvision import transforms
 import gradio as gr
+import spaces
+from gradio_imageslider import ImageSlider
+
+torch.jit.script = lambda f: f
 
 from models.baseline import BiRefNet
 from config import Config
@@ -35,20 +39,22 @@ class ImagePreprocessor():
         return image
 
 
-model = BiRefNet(bb_pretrained=False).to(device)
+model = BiRefNet(bb_pretrained=False)
 state_dict = './BiRefNet_ep580.pth'
 if os.path.exists(state_dict):
-    birefnet_dict = torch.load(state_dict, map_location=device)
+    birefnet_dict = torch.load(state_dict, map_location="cpu")
     unwanted_prefix = '_orig_mod.'
     for k, v in list(birefnet_dict.items()):
         if k.startswith(unwanted_prefix):
             birefnet_dict[k[len(unwanted_prefix):]] = birefnet_dict.pop(k)
     model.load_state_dict(birefnet_dict)
+model = model.to(device)
 model.eval()
 
 
 # def predict(image_1, image_2):
 #     images = [image_1, image_2]
+@spaces.GPU
 def predict(image, resolution='1024x1024'):
     # Image is a RGB numpy array.
     resolution = [int(int(reso)//32*32) for reso in resolution.strip().split('x')]
@@ -74,17 +80,13 @@ def predict(image, resolution='1024x1024'):
         image_preds.append(
             cv2.cvtColor((pred*255).astype(np.uint8), cv2.COLOR_GRAY2RGB)
         )
-    return image_preds[:] if len(images) > 1 else image_preds[0]
 
+    return image, image_preds[0]
 
-examples = [[_] for _ in glob('materials/examples/*')][:]
 
-N = 1
-ipt = [gr.Image() for _ in range(N)]
-opt = [gr.Image() for _ in range(N)]
+examples = [[_] for _ in glob('materials/examples/*')][:]
 
 # Add the option of resolution in a text box.
-ipt += [gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `512x512`. Higher resolutions can be much slower for inference.", label="Resolution")]
 for idx_example, example in enumerate(examples):
     examples[idx_example].append('1024x1024')
 examples.append(examples[-1].copy())
@@ -92,8 +94,8 @@ examples[-1][1] = '512x512'
 
 demo = gr.Interface(
     fn=predict,
-    inputs=ipt,
-    outputs=opt,
+    inputs=['image', gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `512x512`. Higher resolutions can be much slower for inference.", label="Resolution")],
+    outputs=ImageSlider(),
     examples=examples,
     title='Online demo for `Bilateral Reference for High-Resolution Dichotomous Image Segmentation`',
     description=('Upload a picture, our model will give you the binary maps of the highly accurate segmentation of the salient objects in it. :)'

config.py

@@ -1,6 +1,8 @@
 import os
 import math
 
+import torch
+
 
 class Config():
     def __init__(self) -> None:
@@ -97,7 +99,7 @@ class Config():
         self.lambda_adv_d = 3. * (self.lambda_adv_g > 0)
 
         # others
-        self.device = [0, 'cpu'][1]     # .to(0) = .to('cuda:0')
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
 
         self.batch_size_valid = 1
         self.rand_seed = 7

requirements.txt

@@ -1,6 +1,4 @@
---extra-index-url https://download.pytorch.org/whl/cu118
 torch==2.0.1
---extra-index-url https://download.pytorch.org/whl/cu118
 torchvision==0.15.2
 opencv-python==4.9.0.80
 tqdm==4.66.2
@@ -9,3 +7,4 @@ prettytable==3.10.0
 scipy==1.12.0
 scikit-image==0.22.0
 kornia==0.7.1
+gradio_imageslider==0.0.18