66 lines
1.8 KiB
Python
66 lines
1.8 KiB
Python
"""CLI: parse args, load config, run downloads."""
|
||
|
||
import argparse
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
from .config import get_model_tasks, load_config, validate_config
|
||
from .download import run_downloads
|
||
|
||
|
||
def main() -> int:
|
||
parser = argparse.ArgumentParser(
|
||
description="Download ComfyUI models from Hugging Face and Civitai into configured folders."
|
||
)
|
||
parser.add_argument(
|
||
"--config",
|
||
type=Path,
|
||
default=Path("config.yaml"),
|
||
help="Path to YAML config (default: config.yaml in cwd)",
|
||
)
|
||
parser.add_argument(
|
||
"--dry-run",
|
||
action="store_true",
|
||
help="Only print what would be downloaded and where; no writes",
|
||
)
|
||
parser.add_argument(
|
||
"--only",
|
||
nargs="+",
|
||
metavar="TYPE",
|
||
help="Only process these model types (e.g. diffusion_models loras)",
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
try:
|
||
config = load_config(args.config)
|
||
except FileNotFoundError as e:
|
||
print(f"Error: {e}", file=sys.stderr)
|
||
return 1
|
||
except Exception as e:
|
||
print(f"Error loading config: {e}", file=sys.stderr)
|
||
return 1
|
||
|
||
try:
|
||
validate_config(config)
|
||
except ValueError as e:
|
||
print(f"Error: {e}", file=sys.stderr)
|
||
return 1
|
||
|
||
tasks = get_model_tasks(config, only_types=args.only)
|
||
if not tasks:
|
||
print("No model URLs to download. Add URLs under diffusion_models, text_encoders, vaes, upscale_models, or loras in your config.")
|
||
return 0
|
||
|
||
if args.dry_run:
|
||
print("Dry run – no files will be written.\n")
|
||
|
||
ok, fail = run_downloads(config, tasks, dry_run=args.dry_run)
|
||
|
||
if not args.dry_run:
|
||
print(f"\nDone: {ok} succeeded, {fail} failed.")
|
||
return 1 if fail else 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|