69 lines
2.1 KiB
Python
69 lines
2.1 KiB
Python
# Copyright (c) 2025 Salvador E. Tropea
|
|
# Copyright (c) 2025 Instituto Nacional de Tecnología Industrial
|
|
# License: GPLv3
|
|
# Project: ComfyUI-AudioSeparation
|
|
#
|
|
# Tool to download a model from the models Data Base.
|
|
# Run it using: python tool/download_model.py HASH
|
|
import argparse
|
|
import sys
|
|
# Local imports
|
|
import bootstrap # noqa: F401
|
|
from source.db.hash_dir import hash_dir
|
|
from source.db.models_db import load_known_models, get_download_url, cli_add_models_and_db
|
|
from source.utils.downloader import download_model
|
|
from source.utils.logger import main_logger, logger_set_standalone
|
|
from source.utils.misc import cli_add_verbose
|
|
|
|
|
|
def main(args):
|
|
logger_set_standalone(args)
|
|
# Load the JSON metadata file
|
|
model_db = load_known_models(args.json_file)
|
|
if model_db is None:
|
|
main_logger.error("No model database available.")
|
|
sys.exit(3)
|
|
|
|
# Is a valid hash?
|
|
if args.hash not in model_db:
|
|
main_logger.error(f"Nothing known about `{args.hash}`")
|
|
sys.exit(4)
|
|
d = model_db[args.hash]
|
|
|
|
# Check what we have
|
|
downloaded = hash_dir(args.models_dir)
|
|
if args.hash in downloaded:
|
|
main_logger.error(f"`{args.hash}` already downloaded as `{downloaded[args.hash]}`")
|
|
sys.exit(5)
|
|
|
|
# Check we can download it
|
|
url = get_download_url(d)
|
|
if url is None:
|
|
main_logger.error(f"`{args.hash}` can't be downloaded")
|
|
sys.exit(6)
|
|
|
|
# Download the file
|
|
name = d['name']
|
|
try:
|
|
download_model(url, args.models_dir, name, force_urllib=False)
|
|
except Exception as e:
|
|
main_logger.error(f"Failed to download {name} from {url}\n{e}")
|
|
raise
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(
|
|
description="Downloads a model from the database.",
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter
|
|
)
|
|
|
|
# --- File/Path Arguments ---
|
|
parser.add_argument('hash', type=str, help="Hash for the file to download.")
|
|
cli_add_models_and_db(parser)
|
|
|
|
# --- Control Arguments ---
|
|
cli_add_verbose(parser)
|
|
|
|
args = parser.parse_args()
|
|
main(args)
|