fix gfpgan cpu/mps device
This commit is contained in:
parent
8af5df491e
commit
c4968dd0a9
@ -107,7 +107,7 @@ def parse_args():
|
|||||||
help="Enable GFPGAN face restore",
|
help="Enable GFPGAN face restore",
|
||||||
)
|
)
|
||||||
parser.add_argument(
|
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(
|
parser.add_argument(
|
||||||
"--enable-gif",
|
"--enable-gif",
|
||||||
|
@ -17,6 +17,10 @@ class GFPGANPlugin(BasePlugin):
|
|||||||
model_path = download_model(url, model_md5)
|
model_path = download_model(url, model_md5)
|
||||||
logger.info(f"GFPGAN model path: {model_path}")
|
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
|
# Use GFPGAN for face enhancement
|
||||||
self.face_enhancer = MyGFPGANer(
|
self.face_enhancer = MyGFPGANer(
|
||||||
model_path=model_path,
|
model_path=model_path,
|
||||||
|
@ -24,23 +24,27 @@ def test_remove_bg():
|
|||||||
_save(res, "test_remove_bg.png")
|
_save(res, "test_remove_bg.png")
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("device", ["cuda", "cpu"])
|
@pytest.mark.parametrize("device", ["cuda", "cpu", "mps"])
|
||||||
def test_upscale(device):
|
def test_upscale(device):
|
||||||
if device == "cuda" and not torch.cuda.is_available():
|
if device == "cuda" and not torch.cuda.is_available():
|
||||||
return
|
return
|
||||||
|
if device == "mps" and not torch.backends.mps.is_available():
|
||||||
|
return
|
||||||
|
|
||||||
model = RealESRGANUpscaler("realesr-general-x4v3", device)
|
model = RealESRGANUpscaler("realesr-general-x4v3", device)
|
||||||
res = model.forward(bgr_img, 2)
|
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)
|
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):
|
def test_gfpgan(device):
|
||||||
if device == "cuda" and not torch.cuda.is_available():
|
if device == "cuda" and not torch.cuda.is_available():
|
||||||
return
|
return
|
||||||
|
if device == "mps" and not torch.backends.mps.is_available():
|
||||||
|
return
|
||||||
model = GFPGANPlugin(device)
|
model = GFPGANPlugin(device)
|
||||||
res = model(rgb_img, None, None)
|
res = model(rgb_img, None, None)
|
||||||
_save(res, "test_gfpgan.png")
|
_save(res, f"test_gfpgan_{device}.png")
|
||||||
|
Loading…
Reference in New Issue
Block a user