kamp/IOPaint
1
0
Fork 0
Code Issues Pull Requests Releases 1 Activity

add --model-dir

Browse Source
This commit is contained in:
Qing 2023-01-14 21:36:32 +08:00
parent 2d4ca6f3ee
commit 719c7278b4
2 changed files with 20 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
lama_cleaner/__init__.py
View File

@ -2,8 +2,10 @@ import warnings
warnings.simplefilter("ignore", UserWarning) warnings.simplefilter("ignore", UserWarning)
from lama_cleaner.parse_args import parse_args from lama_cleaner.parse_args import parse_args
from lama_cleaner.server import main
def entry_point(): def entry_point():
args = parse_args() 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) main(args)

17
lama_cleaner/parse_args.py
View File

@ -65,10 +65,27 @@ def parse_args():
"--output-dir", type=str, "--output-dir", type=str,
help="Only required when --input is directory. Output directory for all processed images" help="Only required when --input is directory. Output directory for all processed images"
) )
parser.add_argument(
"--model-dir", type=str, default=None,
help="Model download directory (by setting XDG_CACHE_HOME environment variable), "
"by default model downloaded to ~/.cache"
)
parser.add_argument("--disable-model-switch", action="store_true", help="Disable model switch in frontend") parser.add_argument("--disable-model-switch", action="store_true", help="Disable model switch in frontend")
parser.add_argument("--debug", action="store_true") parser.add_argument("--debug", action="store_true")
args = parser.parse_args() args = parser.parse_args()
if args.model_dir is not None:
if os.path.isfile(args.model_dir):
parser.error(f"invalid --model-dir: {args.model_dir} is a file")
if not os.path.exists(args.model_dir):
logger.info(f"Create model cache directory: {args.model_dir}")
Path(args.model_dir).mkdir(exist_ok=True, parents=True)
os.environ["XDG_CACHE_HOME"] = args.model_dir
if args.input is not None: if args.input is not None:
if not os.path.exists(args.input): if not os.path.exists(args.input):
parser.error(f"invalid --input: {args.input} not exists") parser.error(f"invalid --input: {args.input} not exists")