add --model-dir

2023-01-14 21:36:32 +08:00
parent 2d4ca6f3ee
commit 719c7278b4
2 changed files with 20 additions and 1 deletions
--- a/lama_cleaner/init.py
+++ b/lama_cleaner/init.py
@@ -2,8 +2,10 @@ import warnings
 warnings.simplefilter("ignore", UserWarning)

 from lama_cleaner.parse_args import parse_args
-from lama_cleaner.server import main

 def entry_point():
    args = parse_args()
+    # To make os.environ["XDG_CACHE_HOME"] = args.model_cache_dir works for diffusers
+    # https://github.com/huggingface/diffusers/blob/be99201a567c1ccd841dc16fb24e88f7f239c187/src/diffusers/utils/constants.py#L18
+    from lama_cleaner.server import main
    main(args)