autoMate/util/download_weights.py
2025-03-20 14:42:37 +08:00

17 lines
485 B
Python

import os
from pathlib import Path
__WEIGHTS_DIR = Path("weights")
OMNI_PARSER_DIR = os.path.join(__WEIGHTS_DIR, "AI-ModelScope", "OmniParser-v2___0")
def download():
from modelscope import snapshot_download
# Create weights directory
__WEIGHTS_DIR.mkdir(exist_ok=True)
snapshot_download(
'AI-ModelScope/OmniParser-v2.0',
cache_dir='weights',
allow_file_pattern=['icon_detect/model.pt']
)
if __name__ == "__main__":
download()