add model md5 check

2023-02-26 09:19:48 +08:00
parent 64336498ba
commit ecfecac050
9 changed files with 2002 additions and 933 deletions
--- a/lama_cleaner/tests/test_model_md5.py
+++ b/lama_cleaner/tests/test_model_md5.py
@@ -0,0 +1,54 @@
+import os
+import tempfile
+from pathlib import Path
+
+
+def test_load_model():
+    from lama_cleaner.interactive_seg import InteractiveSeg
+    from lama_cleaner.model_manager import ModelManager
+
+    interactive_seg_model = InteractiveSeg()
+
+    models = [
+        "lama",
+        "ldm",
+        "zits",
+        "mat",
+        "fcf",
+        "manga",
+    ]
+    for m in models:
+        ModelManager(
+            name=m,
+            device="cpu",
+            no_half=False,
+            hf_access_token="",
+            disable_nsfw=False,
+            sd_cpu_textencoder=True,
+            sd_run_local=True,
+            local_files_only=True,
+            cpu_offload=True,
+            enable_xformers=False,
+        )
+
+
+# def create_empty_file(tmp_dir, name):
+#     tmp_model_dir = os.path.join(tmp_dir, "torch", "hub", "checkpoints")
+#     Path(tmp_model_dir).mkdir(exist_ok=True, parents=True)
+#     path = os.path.join(tmp_model_dir, name)
+#     with open(path, "w") as f:
+#         f.write("1")
+#
+#
+# def test_load_model_error():
+#     MODELS = [
+#         ("big-lama.pt", "e3aa4aaa15225a33ec84f9f4bc47e500"),
+#         ("cond_stage_model_encode.pt", "23239fc9081956a3e70de56472b3f296"),
+#         ("cond_stage_model_decode.pt", "fe419cd15a750d37a4733589d0d3585c"),
+#         ("diffusion.pt", "b0afda12bf790c03aba2a7431f11d22d"),
+#     ]
+#     with tempfile.TemporaryDirectory() as tmp_dir:
+#         os.environ["XDG_CACHE_HOME"] = tmp_dir
+#         for name, md5 in MODELS:
+#             create_empty_file(tmp_dir, name)
+#             test_load_model()