import os
import gradio as gr
import datetime
import tempfile
from huggingface_hub import hf_hub_download


def download_slow(repo_id):
    os.environ.pop("HF_TRANSFER", None)
    with tempfile.NamedDirectory() as workdir:
        hf_hub_download(
            repo_id,
            filename="pytorch_model.bin",
            force_download=True,
            cache_dir=workdir,
        )


def download_fast(repo_id):
    os.environ["HF_TRANSFER"] = "1"
    with tempfile.NamedDirectory() as workdir:
        hf_hub_download(
            repo_id,
            filename="pytorch_model.bin",
            force_download=True,
            cache_dir=workdir,
        )


def download(repo_id):
    start = datetime.datetime.now()
    download_slow(repo_id)
    taken_slow = datetime.datetime.now() - start

    start = datetime.datetime.now()
    download_fast(repo_id)
    taken_fast = datetime.datetime.now() - start

    return f"""
Slow : {taken_slow}
Fast : {taken_fast}
    """


iface = gr.Interface(fn=download, inputs="text", outputs="text")
iface.launch()