add restoreformer

2023-03-30 16:07:38 +08:00
parent f2e90d3f84
commit c52f733214
9 changed files with 103 additions and 21 deletions
--- a/lama_cleaner/parse_args.py
+++ b/lama_cleaner/parse_args.py
@@ -106,7 +106,16 @@ def parse_args():
    )
    parser.add_argument("--enable-gfpgan", action="store_true", help=GFPGAN_HELP)
    parser.add_argument(
-        "--gfpgan-device", default="cpu", type=str, choices=["cpu", "cuda", "mps"]
+        "--gfpgan-device", default="cpu", type=str, choices=GFPGAN_AVAILABLE_DEVICES
+    )
+    parser.add_argument(
+        "--enable-restoreformer", action="store_true", help=RESTOREFORMER_HELP
+    )
+    parser.add_argument(
+        "--restoreformer-device",
+        default="cpu",
+        type=str,
+        choices=RESTOREFORMER_AVAILABLE_DEVICES,
    )
    parser.add_argument(
        "--enable-gif",