From 10fb00311f313cd578d669401e347ca4acec3af5 Mon Sep 17 00:00:00 2001 From: Qing Date: Tue, 16 Nov 2021 21:21:41 +0800 Subject: [PATCH] fix cache_dir in main.py --- README.md | 12 ++++++++---- main.py | 4 ++-- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 51ca646..effd36f 100644 --- a/README.md +++ b/README.md @@ -19,22 +19,26 @@ You can experience their great online services [here](https://cleanup.pictures/) - Build: `yarn build` ## Docker -Run within a Docker container. Set the `cache_dir` to models location path. + +Run within a Docker container. Set the `CACHE_DIR` to models location path. Optionally add a `-d` option to the `docker run` command below to run as a daemon. ### Build Docker image + ``` docker build -f Dockerfile -t lamacleaner . ``` ### Run Docker (cpu) + ``` -docker run -p 8080:8080 -e cache_dir=/app/models -v $(pwd)/models:/app/models -v $(pwd):/app --rm lamacleaner python3 main.py --device=cpu --port=8080 +docker run -p 8080:8080 -e CACHE_DIR=/app/models -v $(pwd)/models:/app/models -v $(pwd):/app --rm lamacleaner python3 main.py --device=cpu --port=8080 ``` ### Run Docker (gpu) + ``` -docker run --gpus all -p 8080:8080 -e cache_dir=/app/models -v $(pwd)/models:/app/models -v $(pwd):/app --rm lamacleaner python3 main.py --device=cuda --port=8080 +docker run --gpus all -p 8080:8080 -e CACHE_DIR=/app/models -v $(pwd)/models:/app/models -v $(pwd):/app --rm lamacleaner python3 main.py --device=cuda --port=8080 ``` -Then open [http://localhost:8080](http://localhost:8080) \ No newline at end of file +Then open [http://localhost:8080](http://localhost:8080) diff --git a/main.py b/main.py index 780e999..1d41ff3 100644 --- a/main.py +++ b/main.py @@ -26,8 +26,8 @@ os.environ["OPENBLAS_NUM_THREADS"] = NUM_THREADS os.environ["MKL_NUM_THREADS"] = NUM_THREADS os.environ["VECLIB_MAXIMUM_THREADS"] = NUM_THREADS os.environ["NUMEXPR_NUM_THREADS"] = NUM_THREADS -if os.environ["cache_dir"]: - os.environ["TORCH_HOME"] = os.environ["cache_dir"] +if os.environ.get("CACHE_DIR"): + os.environ["TORCH_HOME"] = os.environ["CACHE_DIR"] BUILD_DIR = os.environ.get("LAMA_CLEANER_BUILD_DIR", "./lama_cleaner/app/build")