Init model on GPU

app.py

@@ -1,10 +1,10 @@
+import torch
 import gradio as gr
 import spaces
 from inference_gradio import inference_one_image, model_init
 
 MODEL_PATH = "./checkpoints/docres.pkl"
 
-model = model_init(MODEL_PATH)
 possible_tasks = [
     "dewarping",
     "deshadowing",
@@ -13,14 +13,22 @@ possible_tasks = [
     "binarization",
 ]
 
-@spaces.GPU
+@spaces.GPU(duration=90)
 def run_tasks(image, tasks):
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # load model
+    model = model_init(MODEL_PATH, device)
+
+    # run inference
     bgr_image = image[..., ::-1].copy()
-    bgr_restored_image = inference_one_image(model, bgr_image, tasks)
+    bgr_restored_image = inference_one_image(model, bgr_image, tasks, device)
     if bgr_restored_image.ndim == 3:
         rgb_image = bgr_restored_image[..., ::-1]
     else:
         rgb_image = bgr_restored_image
+
     return rgb_image
 
 

inference_gradio.py

@@ -14,9 +14,6 @@ sys.path.append("./data/MBD/")
 from data.MBD.infer import net1_net2_infer_single_im
 
 
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
-
 def dewarp_prompt(img):
     mask = net1_net2_infer_single_im(img, "data/MBD/checkpoint/mbd.pkl")
     base_coord = utils.getBasecoord(256, 256) / 256
@@ -122,7 +119,7 @@ def binarization_promptv2(img):
     )
 
 
-def dewarping(model, im_org):
+def dewarping(model, im_org, device):
     INPUT_SIZE = 256
     im_masked, prompt_org = dewarp_prompt(im_org.copy())
 
@@ -131,10 +128,10 @@ def dewarping(model, im_org):
     im_masked = cv2.resize(im_masked, (INPUT_SIZE, INPUT_SIZE))
     im_masked = im_masked / 255.0
     im_masked = torch.from_numpy(im_masked.transpose(2, 0, 1)).unsqueeze(0)
-    im_masked = im_masked.float().to(DEVICE)
+    im_masked = im_masked.float().to(device)
 
     prompt = torch.from_numpy(prompt_org.transpose(2, 0, 1)).unsqueeze(0)
-    prompt = prompt.float().to(DEVICE)
+    prompt = prompt.float().to(device)
 
     in_im = torch.cat((im_masked, prompt), dim=1)
 
@@ -158,7 +155,7 @@ def dewarping(model, im_org):
     return prompt_org[:, :, 0], prompt_org[:, :, 1], prompt_org[:, :, 2], out_im
 
 
-def appearance(model, im_org):
+def appearance(model, im_org, device):
     MAX_SIZE = 1600
     # obtain im and prompt
     h, w = im_org.shape[:2]
@@ -176,7 +173,7 @@ def appearance(model, im_org):
     in_im = torch.from_numpy(in_im.transpose(2, 0, 1)).unsqueeze(0)
 
     # inference
-    in_im = in_im.half().to(DEVICE)
+    in_im = in_im.half().to(device)
     model = model.half()
     with torch.no_grad():
         pred = model(in_im)
@@ -198,7 +195,7 @@ def appearance(model, im_org):
     return prompt[:, :, 0], prompt[:, :, 1], prompt[:, :, 2], out_im
 
 
-def deshadowing(model, im_org):
+def deshadowing(model, im_org, device):
     MAX_SIZE = 1600
     # obtain im and prompt
     h, w = im_org.shape[:2]
@@ -216,7 +213,7 @@ def deshadowing(model, im_org):
     in_im = torch.from_numpy(in_im.transpose(2, 0, 1)).unsqueeze(0)
 
     # inference
-    in_im = in_im.half().to(DEVICE)
+    in_im = in_im.half().to(device)
     model = model.half()
     with torch.no_grad():
         pred = model(in_im)
@@ -238,16 +235,16 @@ def deshadowing(model, im_org):
     return prompt[:, :, 0], prompt[:, :, 1], prompt[:, :, 2], out_im
 
 
-def deblurring(model, im_org):
+def deblurring(model, im_org, device):
     # setup image
     in_im, padding_h, padding_w = stride_integral(im_org, 8)
     prompt = deblur_prompt(in_im)
     in_im = np.concatenate((in_im, prompt), -1)
     in_im = in_im / 255.0
     in_im = torch.from_numpy(in_im.transpose(2, 0, 1)).unsqueeze(0)
-    in_im = in_im.half().to(DEVICE)
+    in_im = in_im.half().to(device)
     # inference
-    model.to(DEVICE)
+    model.to(device)
     model.eval()
     model = model.half()
     with torch.no_grad():
@@ -260,7 +257,7 @@ def deblurring(model, im_org):
     return prompt[:, :, 0], prompt[:, :, 1], prompt[:, :, 2], out_im
 
 
-def binarization(model, im_org):
+def binarization(model, im_org, device):
     im, padding_h, padding_w = stride_integral(im_org, 8)
     prompt = binarization_promptv2(im)
     h, w = im.shape[:2]
@@ -268,7 +265,7 @@ def binarization(model, im_org):
 
     in_im = in_im / 255.0
     in_im = torch.from_numpy(in_im.transpose(2, 0, 1)).unsqueeze(0)
-    in_im = in_im.to(DEVICE)
+    in_im = in_im.to(device)
     model = model.half()
     in_im = in_im.half()
     with torch.no_grad():
@@ -283,7 +280,7 @@ def binarization(model, im_org):
     return prompt[:, :, 0], prompt[:, :, 1], prompt[:, :, 2], out_im
 
 
-def model_init(model_path):
+def model_init(model_path, device):
     # prepare model
     model = restormer_arch.Restormer(
         inp_channels=6,
@@ -298,7 +295,7 @@ def model_init(model_path):
         dual_pixel_task=True,
     )
 
-    if DEVICE == "cpu":
+    if device == "cpu":
         state = convert_state_dict(
             torch.load(model_path, map_location="cpu")["model_state"]
         )
@@ -309,7 +306,7 @@ def model_init(model_path):
     model.load_state_dict(state)
 
     model.eval()
-    model = model.to(DEVICE)
+    model = model.to(device)
     return model
 
 
@@ -328,11 +325,11 @@ def resize(image, max_size):
     return image
 
 
-def inference_one_image(model, image, tasks):
+def inference_one_image(model, image, tasks, device):
     # image should be in BGR format
 
     if "dewarping" in tasks:
-        *_, image = dewarping(model, image)
+        *_, image = dewarping(model, image, device)
     
     # if only dewarping return here
     if len(tasks) == 1 and "dewarping" in tasks:
@@ -341,12 +338,12 @@ def inference_one_image(model, image, tasks):
     image = resize(image, 1536)
 
     if "deshadowing" in tasks:
-        *_, image = deshadowing(model, image)
+        *_, image = deshadowing(model, image, device)
     if "appearance" in tasks:
-        *_, image = appearance(model, image)
+        *_, image = appearance(model, image, device)
     if "deblurring" in tasks:
-        *_, image = deblurring(model, image)
+        *_, image = deblurring(model, image, device)
     if "binarization" in tasks:
-        *_, image = binarization(model, image)
+        *_, image = binarization(model, image, device)
     
     return image