import requests
import os
from huggingface_hub import hf_hub_download
from utils import custom_drive_cache_dir, get_drive

HF_TOKEN = os.getenv("HF_TOKEN")

ANIME2SKETCH_MODEL = {"REPO_ID": "p1atdev/Anime2Sketch", "FILENAME": "netG.pth"}

def download_anime2sketch_model():
    if os.path.exists("./models/netG.pth"):
        return

    drive = get_drive("./models/netG.pth")
    with custom_drive_cache_dir(drive) as cache_dir:
        hf_hub_download(
            repo_id=ANIME2SKETCH_MODEL["REPO_ID"],
            filename=ANIME2SKETCH_MODEL["FILENAME"],
            local_dir="./models",
            use_auth_token=HF_TOKEN,
            local_dir_use_symlinks=False,
            cache_dir=cache_dir,
        )


def setup():
    if not os.path.exists("./models"):
        os.makedirs("./models")
    download_anime2sketch_model()