add global try.

app.py

@@ -103,98 +103,101 @@ os.makedirs('output', exist_ok=True)
 
 def inference(image, background_enhance, face_upsample, upscale, codeformer_fidelity):
     """Run a single prediction on the model"""
-    # take the default setting for the demo
-    has_aligned = False
-    only_center_face = False
-    draw_box = False
-    detection_model = "retinaface_resnet50"
-
-    upscale = int(upscale) # covert type to int
-    face_helper = FaceRestoreHelper(
-        upscale,
-        face_size=512,
-        crop_ratio=(1, 1),
-        det_model=detection_model,
-        save_ext="png",
-        use_parse=True,
-        device=device,
-    )
-    bg_upsampler = upsampler if background_enhance else None
-    face_upsampler = upsampler if face_upsample else None
-
-    img = cv2.imread(str(image), cv2.IMREAD_COLOR)
-
-    if has_aligned:
-        # the input faces are already cropped and aligned
-        img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
-        face_helper.is_gray = is_gray(img, threshold=5)
-        if face_helper.is_gray:
-            print('Grayscale input: True')
-        face_helper.cropped_faces = [img]
-    else:
-        face_helper.read_image(img)
-        # get face landmarks for each face
-        num_det_faces = face_helper.get_face_landmarks_5(
-          only_center_face=only_center_face, resize=640, eye_dist_threshold=5
-        )
-        print(f"\tdetect {num_det_faces} faces")
-        # align and warp each face
-        face_helper.align_warp_face()
-
-    # face restoration for each cropped face
-    for idx, cropped_face in enumerate(face_helper.cropped_faces):
-        # prepare data
-        cropped_face_t = img2tensor(
-            cropped_face / 255.0, bgr2rgb=True, float32=True
+    try: # global try
+        # take the default setting for the demo
+        has_aligned = False
+        only_center_face = False
+        draw_box = False
+        detection_model = "retinaface_resnet50"
+
+        upscale = int(upscale) # covert type to int
+        face_helper = FaceRestoreHelper(
+            upscale,
+            face_size=512,
+            crop_ratio=(1, 1),
+            det_model=detection_model,
+            save_ext="png",
+            use_parse=True,
+            device=device,
         )
-        normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
-        cropped_face_t = cropped_face_t.unsqueeze(0).to(device)
-
-        try:
-            with torch.no_grad():
-                output = codeformer_net(
-                    cropped_face_t, w=codeformer_fidelity, adain=True
-                )[0]
-                restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
-            del output
-            torch.cuda.empty_cache()
-        except Exception as error:
-            print(f"\tFailed inference for CodeFormer: {error}")
-            restored_face = tensor2img(
-                cropped_face_t, rgb2bgr=True, min_max=(-1, 1)
-            )
-
-        restored_face = restored_face.astype("uint8")
-        face_helper.add_restored_face(restored_face)
-
-    # paste_back
-    if not has_aligned:
-        # upsample the background
-        if bg_upsampler is not None:
-            # Now only support RealESRGAN for upsampling background
-            bg_img = bg_upsampler.enhance(img, outscale=upscale)[0]
+        bg_upsampler = upsampler if background_enhance else None
+        face_upsampler = upsampler if face_upsample else None
+
+        img = cv2.imread(str(image), cv2.IMREAD_COLOR)
+
+        if has_aligned:
+            # the input faces are already cropped and aligned
+            img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
+            face_helper.is_gray = is_gray(img, threshold=5)
+            if face_helper.is_gray:
+                print('Grayscale input: True')
+            face_helper.cropped_faces = [img]
         else:
-            bg_img = None
-        face_helper.get_inverse_affine(None)
-        # paste each restored face to the input image
-        if face_upsample and face_upsampler is not None:
-            restored_img = face_helper.paste_faces_to_input_image(
-                upsample_img=bg_img,
-                draw_box=draw_box,
-                face_upsampler=face_upsampler,
+            face_helper.read_image(img)
+            # get face landmarks for each face
+            num_det_faces = face_helper.get_face_landmarks_5(
+            only_center_face=only_center_face, resize=640, eye_dist_threshold=5
             )
-        else:
-            restored_img = face_helper.paste_faces_to_input_image(
-                upsample_img=bg_img, draw_box=draw_box
+            print(f"\tdetect {num_det_faces} faces")
+            # align and warp each face
+            face_helper.align_warp_face()
+
+        # face restoration for each cropped face
+        for idx, cropped_face in enumerate(face_helper.cropped_faces):
+            # prepare data
+            cropped_face_t = img2tensor(
+                cropped_face / 255.0, bgr2rgb=True, float32=True
             )
-
-    # save restored img
-    save_path = f'output/out.png'
-    imwrite(restored_img, str(save_path))
-
-    restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)
-    return restored_img, save_path
-
+            normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
+            cropped_face_t = cropped_face_t.unsqueeze(0).to(device)
+
+            try:
+                with torch.no_grad():
+                    output = codeformer_net(
+                        cropped_face_t, w=codeformer_fidelity, adain=True
+                    )[0]
+                    restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
+                del output
+                torch.cuda.empty_cache()
+            except Exception as error:
+                print(f"\tFailed inference for CodeFormer: {error}")
+                restored_face = tensor2img(
+                    cropped_face_t, rgb2bgr=True, min_max=(-1, 1)
+                )
+
+            restored_face = restored_face.astype("uint8")
+            face_helper.add_restored_face(restored_face)
+
+        # paste_back
+        if not has_aligned:
+            # upsample the background
+            if bg_upsampler is not None:
+                # Now only support RealESRGAN for upsampling background
+                bg_img = bg_upsampler.enhance(img, outscale=upscale)[0]
+            else:
+                bg_img = None
+            face_helper.get_inverse_affine(None)
+            # paste each restored face to the input image
+            if face_upsample and face_upsampler is not None:
+                restored_img = face_helper.paste_faces_to_input_image(
+                    upsample_img=bg_img,
+                    draw_box=draw_box,
+                    face_upsampler=face_upsampler,
+                )
+            else:
+                restored_img = face_helper.paste_faces_to_input_image(
+                    upsample_img=bg_img, draw_box=draw_box
+                )
+
+        # save restored img
+        save_path = f'output/out.png'
+        imwrite(restored_img, str(save_path))
+
+        restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)
+        return restored_img, save_path
+    except Exception as error:
+        print('global exception', error)
+        return None, None
 
 
 title = "CodeFormer: Robust Face Restoration and Enhancement Network"
@@ -230,7 +233,7 @@ Redistribution and use for non-commercial purposes should follow this license.
 
 If you have any questions, please feel free to reach me out at <b>shangchenzhou@gmail.com</b>.
 
-![visitors](https://visitor-badge.glitch.me/badge?page_id=sczhou/CodeFormer)
+![visitors](https://visitor-badge.laobi.icu/badge?page_id=sczhou/CodeFormer)
 """
 
 demo = gr.Interface(