kamp/IOPaint
1
0
Fork 0
Code Issues Pull Requests Releases 1 Activity

fix gfpgan cpu/mps device

Browse Source
This commit is contained in:
Qing 2023-03-28 16:36:41 +08:00
parent 8af5df491e
commit c4968dd0a9
3 changed files with 14 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
lama_cleaner/parse_args.py
View File

@ -107,7 +107,7 @@ def parse_args():
help="Enable GFPGAN face restore",
)
parser.add_argument(
"--gfpgan-device", default="cpu", type=str, choices=["cpu", "cuda"]
"--gfpgan-device", default="cpu", type=str, choices=["cpu", "cuda", "mps"]
)
parser.add_argument(
"--enable-gif",

4
lama_cleaner/plugins/gfpgan_plugin.py
View File

@ -17,6 +17,10 @@ class GFPGANPlugin(BasePlugin):
model_path = download_model(url, model_md5)
logger.info(f"GFPGAN model path: {model_path}")
import facexlib
if hasattr(facexlib.detection.retinaface, "device"):
facexlib.detection.retinaface.device = device
# Use GFPGAN for face enhancement
self.face_enhancer = MyGFPGANer(
model_path=model_path,

14
lama_cleaner/tests/test_plugins.py
View File

@ -24,23 +24,27 @@ def test_remove_bg():
_save(res, "test_remove_bg.png")
@pytest.mark.parametrize("device", ["cuda", "cpu"])
@pytest.mark.parametrize("device", ["cuda", "cpu", "mps"])
def test_upscale(device):
if device == "cuda" and not torch.cuda.is_available():
return
if device == "mps" and not torch.backends.mps.is_available():
return
model = RealESRGANUpscaler("realesr-general-x4v3", device)
res = model.forward(bgr_img, 2)
_save(res, "test_upscale_x2.png")
_save(res, f"test_upscale_x2_{device}.png")
res = model.forward(bgr_img, 4)
_save(res, "test_upscale_x4.png")
_save(res, f"test_upscale_x4_{device}.png")
@pytest.mark.parametrize("device", ["cuda", "cpu"])
@pytest.mark.parametrize("device", ["cuda", "cpu", "mps"])
def test_gfpgan(device):
if device == "cuda" and not torch.cuda.is_available():
return
if device == "mps" and not torch.backends.mps.is_available():
return
model = GFPGANPlugin(device)
res = model(rgb_img, None, None)
_save(res, "test_gfpgan.png")
_save(res, f"test_gfpgan_{device}.png")