from huggingface_hub import hf_hub_download, snapshot_download
import transformers

MOUNT_PATH = "/suno/models"


def download_models():
    ignore = ["*.bin", "*.onnx_data", "*/diffusion_pytorch_model.safetensors"]
    hf_download_path = hf_hub_download(
        "ByteDance/SDXL-Lightning", "sdxl_lightning_4step_unet.safetensors", local_dir=MOUNT_PATH
    )
    print("downloaded to", hf_download_path)
    snapshot_download_path = snapshot_download(
        "stabilityai/stable-diffusion-xl-base-1.0", ignore_patterns=ignore
    )
    print("downloaded to", snapshot_download_path)
    print("finish downloads")

    import transformers

    transformers.utils.move_cache()