Merge pull request #829 from RishonLi/patch-1

Update face_enhancer.py for apple silicon mps
pull/844/head
Kenneth Estanislao 2024-12-16 22:30:50 +08:00 committed by GitHub
commit 5f5fe8890a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 1 deletions

View File

@ -49,7 +49,10 @@ def get_face_enhancer() -> Any:
with THREAD_LOCK:
if FACE_ENHANCER is None:
model_path = os.path.join(models_dir, 'GFPGANv1.4.pth')
FACE_ENHANCER = gfpgan.GFPGANer(model_path=model_path, upscale=1) # type: ignore[attr-defined]
mps_device = None
if torch.backends.mps.is_available():
mps_device = torch.device("mps")
FACE_ENHANCER = gfpgan.GFPGANer(model_path=model_path, upscale=1, device=mps_device) # type: ignore[attr-defined]
return FACE_ENHANCER