download_models() {
    ED_MODEL_URL="https://huggingface.co/Deressa/GenConViT/resolve/main/genconvit_ed_inference.pth"
    # VAE_MODEL_URL="https://huggingface.co/Deressa/GenConViT/resolve/main/genconvit_vae_inference.pth"

    ED_MODEL_PATH="./pretrained_models/genconvit_ed_inference.pth"
    # VAE_MODEL_PATH="./pretrained_models/genconvit_vae_inference.pth"

    mkdir -p pretrained_models

    if [ ! -f "$ED_MODEL_PATH" ]; then
        wget -P ./pretrained_models "$ED_MODEL_URL"
    fi

    # if [ ! -f "$VAE_MODEL_PATH" ]; then
    #     wget -P ./pretrained_models "$VAE_MODEL_URL"
    # fi
}


download_models