| from __future__ import annotations |
|
|
| import logging |
|
|
| import torch |
|
|
| from modules import ( |
| devices, |
| errors, |
| face_restoration, |
| face_restoration_utils, |
| modelloader, |
| 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' |
|
|
| |
| codeformer: face_restoration.FaceRestoration | None = None |
|
|
|
|
| class FaceRestorerCodeFormer(face_restoration_utils.CommonFaceRestoration): |
| def name(self): |
| return "CodeFormer" |
|
|
| def load_net(self) -> torch.Module: |
| for model_path in modelloader.load_models( |
| model_path=self.model_path, |
| model_url=model_url, |
| command_path=self.model_path, |
| download_name=model_download_name, |
| ext_filter=['.pth'], |
| ): |
| return modelloader.load_spandrel_model( |
| model_path, |
| device=devices.device_codeformer, |
| expected_architecture='CodeFormer', |
| ).model |
| raise ValueError("No codeformer model found") |
|
|
| def get_device(self): |
| return devices.device_codeformer |
|
|
| def restore(self, np_image, w: float | None = None): |
| if w is None: |
| w = getattr(shared.opts, "code_former_weight", 0.5) |
|
|
| def restore_face(cropped_face_t): |
| assert self.net is not None |
| return self.net(cropped_face_t, w=w, adain=True)[0] |
|
|
| return self.restore_with_helper(np_image, restore_face) |
|
|
|
|
| def setup_model(dirname: str) -> None: |
| global codeformer |
| try: |
| codeformer = FaceRestorerCodeFormer(dirname) |
| shared.face_restorers.append(codeformer) |
| except Exception: |
| errors.report("Error setting up CodeFormer", exc_info=True) |
|
|
