Progress when downloading with huggingface_hub.snapshot_download

Using Streamlit
1

Does anyone have working progress bars when downloading a model repo with huggingface_hub? It’s a pretty common requirement but can’t find any examples. I’ve been going down the monkey patching of tqdm route and making progress on that, but it’s pretty hacky, as it needs to create multiple progress bars for each file.

2

This is showing up on the top of Google now, so answering my own question here:

from typing import Callable, NoReturn

import streamlit as st
from tqdm.auto import tqdm

def patch_tqdm(progress_bar: Callable[..., NoReturn]):
    """Monkey patches tqdm to update the progress bar in Streamlit when
    downloading models through hf_hub_download."""
    original_update = tqdm.update

    def patched_update(self, n=1):
        if self.n is not None and self.total is not None:
            # Format the information
            downloaded = self.n / 1024 / 1024
            total_size = self.total / 1024 / 1024
            speed = (
                self.format_dict["rate"] / 1024 / 1024
                if self.format_dict["rate"]
                else 0.001
            )
            time_left = format_time((total_size - downloaded) / speed)
            progress_bar.progress(
                self.n / self.total,
                f"Downloaded: {int(downloaded)}/{int(total_size)} MB Speed: {speed:.2f} MB/s Remaining: {time_left}",
            )

        return original_update(self, n)

    def __del__(self):
        progress_bar.empty()
        self.close()

    tqdm.update = patched_update
    tqdm.__del__ = __del__

You can now create a progress bar with progress_bar = st.empty() wherever you like, and pass it into patch_tqdm to hook in from tqdm > streamlit. Now when you hf_hub_download the progress will show up in streamlit, and when the download finishes the bar will disappear. Alternatively you can change __del__ to progress_bar.progress(1.0, "Download complete!") if you want it to stick around.

It’s a little hacky, will see if hugginface_hub will accept a contribution under extras for a class that inherits from tqdm so you can pip install huggingface_hub[streamlit] to get this functionality.