from huggingface_hub import hf_hub_download

# Define the files to download
files = [
    "icon_detect/train_args.yaml",
    "icon_detect/model.pt",
    "icon_detect/model.yaml",
    "icon_caption/config.json",
    "icon_caption/generation_config.json",
    "icon_caption/model.safetensors"
]

# Download each file to the 'weights' folder
for file in files:
    hf_hub_download(repo_id="microsoft/OmniParser-v2.0", filename=file, local_dir="weights")