face

modules/codeformer_model.py

@@ -1,8 +1,7 @@
 from __future__ import annotations
 
-import logging
-
 import torch
+import os
 
 from modules import (
     devices,
@@ -13,12 +12,6 @@ from modules import (
     shared,
 )
 
-logger = logging.getLogger(__name__)
-
-model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
-model_download_name = 'codeformer-v0.1.0.pth'
-
-# used by e.g. postprocessing_codeformer.py
 codeformer: face_restoration.FaceRestoration | None = None
 
 
@@ -26,23 +19,21 @@ class FaceRestorerCodeFormer(face_restoration_utils.CommonFaceRestoration):
     def name(self):
         return "CodeFormer"
 
-    def load_net(self) -> torch.Module:
+    def load_net(self) -> torch.nn.Module:
+        os.makedirs(self.model_path, exist_ok=True)
         for model_path in modelloader.load_models(
             model_path=self.model_path,
-            model_url=model_url,
+            model_url="https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth",
             command_path=self.model_path,
-            download_name=model_download_name,
-            ext_filter=['.pth'],
+            download_name="codeformer-v0.1.0.pth",
+            ext_filter=[".pth"],
         ):
             return modelloader.load_spandrel_model(
                 model_path,
                 device=devices.device_codeformer,
-                expected_architecture='CodeFormer',
+                expected_architecture="CodeFormer",
             ).model
-        raise ValueError("No codeformer model found")
-
-    def get_device(self):
-        return devices.device_codeformer
+        raise ValueError("No CodeFormer Model Found")
 
     def restore(self, np_image, w: float | None = None):
         if w is None:
@@ -50,7 +41,7 @@ class FaceRestorerCodeFormer(face_restoration_utils.CommonFaceRestoration):
 
         def restore_face(cropped_face_t):
             assert self.net is not None
-            return self.net(cropped_face_t, weight=w, adain=True)[0]
+            return self.net(cropped_face_t, w=w, adain=True)[0]
 
         return self.restore_with_helper(np_image, restore_face)
 

modules/face_restoration_utils.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import logging
-import os
 from functools import cached_property
 from typing import TYPE_CHECKING, Callable
 
@@ -43,14 +42,15 @@ def rgb_tensor_to_bgr_image(tensor: torch.Tensor, *, min_max=(0.0, 1.0)) -> np.n
 def create_face_helper(device) -> FaceRestoreHelper:
     from facexlib.detection import retinaface
     from facexlib.utils.face_restoration_helper import FaceRestoreHelper
-    if hasattr(retinaface, 'device'):
+
+    if hasattr(retinaface, "device"):
         retinaface.device = device
     return FaceRestoreHelper(
         upscale_factor=1,
         face_size=512,
         crop_ratio=(1, 1),
-        det_model='retinaface_resnet50',
-        save_ext='png',
+        det_model="retinaface_resnet50",
+        save_ext="png",
         use_parse=True,
         device=device,
     )
@@ -67,6 +67,7 @@ def restore_with_face_helper(
     `restore_face` should take a cropped face image and return a restored face image.
     """
     from torchvision.transforms.functional import normalize
+
     np_image = np_image[:, :, ::-1]
     original_resolution = np_image.shape[0:2]
 
@@ -87,10 +88,10 @@ def restore_with_face_helper(
                     cropped_face_t = restore_face(cropped_face_t)
                 devices.torch_gc()
             except Exception:
-                errors.report('Failed face-restoration inference', exc_info=True)
+                errors.report("Failed face-restoration inference", exc_info=True)
 
             restored_face = rgb_tensor_to_bgr_image(cropped_face_t, min_max=(-1, 1))
-            restored_face = (restored_face * 255.0).astype('uint8')
+            restored_face = (restored_face * 255.0).astype("uint8")
             face_helper.add_restored_face(restored_face)
 
         logger.debug("Merging restored faces into image")
@@ -112,7 +113,7 @@ def restore_with_face_helper(
 
 
 class CommonFaceRestoration(face_restoration.FaceRestoration):
-    net: torch.Module | None
+    net: torch.nn.Module | None
     model_url: str
     model_download_name: str
 
@@ -120,7 +121,6 @@ class CommonFaceRestoration(face_restoration.FaceRestoration):
         super().__init__()
         self.net = None
         self.model_path = model_path
-        os.makedirs(model_path, exist_ok=True)
 
     @cached_property
     def face_helper(self) -> FaceRestoreHelper:
@@ -136,10 +136,10 @@ class CommonFaceRestoration(face_restoration.FaceRestoration):
             self.face_helper.face_parse.to(device)
 
     def get_device(self):
-        raise NotImplementedError("get_device must be implemented by subclasses")
+        return devices.get_optimal_device()
 
-    def load_net(self) -> torch.Module:
-        raise NotImplementedError("load_net must be implemented by subclasses")
+    def load_net(self) -> torch.nn.Module:
+        raise NotImplementedError
 
     def restore_with_helper(
         self,

modules/gfpgan_model.py

@@ -1,9 +1,7 @@
 from __future__ import annotations
 
-import logging
-import os
-
 import torch
+import os
 
 from modules import (
     devices,
@@ -14,9 +12,6 @@ from modules import (
     shared,
 )
 
-logger = logging.getLogger(__name__)
-model_url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
-model_download_name = "GFPGANv1.4.pth"
 gfpgan_face_restorer: face_restoration.FaceRestoration | None = None
 
 
@@ -24,35 +19,21 @@ class FaceRestorerGFPGAN(face_restoration_utils.CommonFaceRestoration):
     def name(self):
         return "GFPGAN"
 
-    def get_device(self):
-        return devices.device_gfpgan
-
-    def load_net(self) -> torch.Module:
+    def load_net(self) -> torch.nn.Module:
+        os.makedirs(self.model_path, exist_ok=True)
         for model_path in modelloader.load_models(
             model_path=self.model_path,
-            model_url=model_url,
+            model_url="https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth",
             command_path=self.model_path,
-            download_name=model_download_name,
-            ext_filter=['.pth'],
+            download_name="GFPGANv1.4.pth",
+            ext_filter=[".pth"],
         ):
-            if 'GFPGAN' in os.path.basename(model_path):
-                return modelloader.load_spandrel_model(
-                    model_path,
-                    device=self.get_device(),
-                    expected_architecture='GFPGAN',
-                ).model
-
-        #   if reach here, model not found. previous code will download it iff there are no models in GFPGAN directory
-        #   this will download it if the supporting models exist
-        try:
-            GFPGANmodel = modelloader.load_file_from_url(model_url, model_dir=self.model_path, file_name=model_download_name)
             return modelloader.load_spandrel_model(
-                GFPGANmodel,
-                device=self.get_device(),
-                expected_architecture='GFPGAN',
+                model_path,
+                device=devices.device_gfpgan,
+                expected_architecture="GFPGAN",
             ).model
-        except:
-            raise ValueError("No GFPGAN model found")
+        raise ValueError("No GFPGAN Model Found")
 
     def restore(self, np_image):
         def restore_face(cropped_face_t):
@@ -65,7 +46,7 @@ class FaceRestorerGFPGAN(face_restoration_utils.CommonFaceRestoration):
 def gfpgan_fix_faces(np_image):
     if gfpgan_face_restorer:
         return gfpgan_face_restorer.restore(np_image)
-    logger.warning("GFPGAN face restorer not set up")
+    print("WARNING: GFPGAN face restorer was not set up")
     return np_image