Merge branch 'dev' into patch-1

1 files changed, 7 insertions, 18 deletions

diff --git a/modules/codeformer_model.py b/modules/codeformer_model.py
index 8d84bbc9..f293acf5 100644
--- a/modules/codeformer_model.py
+++ b/modules/codeformer_model.py
@@ -1,13 +1,11 @@
 import os

-import sys

-import traceback

 

 import cv2

 import torch

 

 import modules.face_restoration

 import modules.shared

-from modules import shared, devices, modelloader

+from modules import shared, devices, modelloader, errors

 from modules.paths import models_path

 

 # codeformer people made a choice to include modified basicsr library to their project which makes

@@ -17,14 +15,11 @@ model_dir = "Codeformer"
 model_path = os.path.join(models_path, model_dir)

 model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'

 

-have_codeformer = False

 codeformer = None

 

 

 def setup_model(dirname):

-    global model_path

-    if not os.path.exists(model_path):

-        os.makedirs(model_path)

+    os.makedirs(model_path, exist_ok=True)

 

     path = modules.paths.paths.get("CodeFormer", None)

     if path is None:

@@ -33,11 +28,9 @@ def setup_model(dirname):
     try:

         from torchvision.transforms.functional import normalize

         from modules.codeformer.codeformer_arch import CodeFormer

-        from basicsr.utils.download_util import load_file_from_url

-        from basicsr.utils import imwrite, img2tensor, tensor2img

+        from basicsr.utils import img2tensor, tensor2img

         from facelib.utils.face_restoration_helper import FaceRestoreHelper

         from facelib.detection.retinaface import retinaface

-        from modules.shared import cmd_opts

 

         net_class = CodeFormer

 

@@ -96,7 +89,7 @@ def setup_model(dirname):
                 self.face_helper.get_face_landmarks_5(only_center_face=False, resize=640, eye_dist_threshold=5)

                 self.face_helper.align_warp_face()

 

-                for idx, cropped_face in enumerate(self.face_helper.cropped_faces):

+                for cropped_face in self.face_helper.cropped_faces:

                     cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)

                     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(devices.device_codeformer)

@@ -107,8 +100,8 @@ def setup_model(dirname):
                             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}', file=sys.stderr)

+                    except Exception:

+                        errors.report('Failed inference for CodeFormer', exc_info=True)

                         restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1))

 

                     restored_face = restored_face.astype('uint8')

@@ -129,15 +122,11 @@ def setup_model(dirname):
 

                 return restored_img

 

-        global have_codeformer

-        have_codeformer = True

-

         global codeformer

         codeformer = FaceRestorerCodeFormer(dirname)

         shared.face_restorers.append(codeformer)

 

     except Exception:

-        print("Error setting up CodeFormer:", file=sys.stderr)

-        print(traceback.format_exc(), file=sys.stderr)

+        errors.report("Error setting up CodeFormer", exc_info=True)

 

    # sys.path = stored_sys_path