import os
from huggingface_hub import hf_hub_download
import shutil

def download_models():
    # Create directories if they don't exist
    os.makedirs("weights/icon_detect", exist_ok=True)
    os.makedirs("weights/icon_caption_florence", exist_ok=True)

    # Define file mappings (repository path -> local path)
    files_to_download = {
        "icon_caption_florence/config.json": "weights/icon_caption_florence/config.json",
        "icon_caption_florence/generation_config.json": "weights/icon_caption_florence/generation_config.json",
        "icon_caption_florence/model.safetensors": "weights/icon_caption_florence/model.safetensors",
        "icon_detect/best.pt": "weights/icon_detect/best.pt"
    }

    # Download each file
    for repo_path, local_path in files_to_download.items():
        if not os.path.exists(local_path):
            print(f"Downloading {repo_path}...")
            try:
                downloaded_file = hf_hub_download(
                    repo_id="banao-tech/OmniParser",
                    filename=repo_path,
                    local_dir="temp"
                )
                # Move the file to the correct location
                os.makedirs(os.path.dirname(local_path), exist_ok=True)
                shutil.move(downloaded_file, local_path)
                print(f"Successfully downloaded and moved to {local_path}")
            except Exception as e:
                print(f"Error downloading {repo_path}: {str(e)}")

    # Clean up temp directory
    if os.path.exists("temp"):
        shutil.rmtree("temp")

if __name__ == "__main__":
    download_models()