From ba8086ace290bb97381de44a6ca7dbcae949bf73 Mon Sep 17 00:00:00 2001 From: Haoming Date: Thu, 31 Jul 2025 14:38:04 +0800 Subject: [PATCH] face --- modules/codeformer_model.py | 27 +++++++------------- modules/face_restoration_utils.py | 22 ++++++++--------- modules/gfpgan_model.py | 41 +++++++++---------------------- 3 files changed, 31 insertions(+), 59 deletions(-) diff --git a/modules/codeformer_model.py b/modules/codeformer_model.py index 0b353353..d8b9d875 100644 --- a/modules/codeformer_model.py +++ b/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) diff --git a/modules/face_restoration_utils.py b/modules/face_restoration_utils.py index 210e2dff..611c5ff3 100644 --- a/modules/face_restoration_utils.py +++ b/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, diff --git a/modules/gfpgan_model.py b/modules/gfpgan_model.py index dffe1613..1236b6f8 100644 --- a/modules/gfpgan_model.py +++ b/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