add download command

This commit is contained in:
Qing
2023-11-16 21:12:06 +08:00
parent 20e660aa4a
commit 1d145d1cd6
17 changed files with 233 additions and 67 deletions

24
lama_cleaner/download.py Normal file
View File

@@ -0,0 +1,24 @@
import os
from loguru import logger
from pathlib import Path
def cli_download_model(model: str, model_dir: str):
if os.path.isfile(model_dir):
raise ValueError(f"invalid --model-dir: {model_dir} is a file")
if not os.path.exists(model_dir):
logger.info(f"Create model cache directory: {model_dir}")
Path(model_dir).mkdir(exist_ok=True, parents=True)
os.environ["XDG_CACHE_HOME"] = model_dir
from lama_cleaner.model_manager import models
if model in models:
logger.info(f"Downloading {model}...")
models[model].download()
logger.info(f"Done.")
else:
logger.info(f"Downloading model from Huggingface: {model}")