mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
52 lines
1.5 KiB
Python
52 lines
1.5 KiB
Python
|
from pathlib import Path
|
||
|
|
from urllib import request
|
||
|
|
|
||
|
|
from tqdm import tqdm
|
||
|
|
|
||
|
|
from invokeai.backend.util.logging import InvokeAILogger
|
||
|
|
|
||
|
|
|
||
|
|
class ProgressBar:
|
||
|
|
"""Simple progress bar for urllib.request.urlretrieve using tqdm."""
|
||
|
|
|
||
|
|
def __init__(self, model_name: str = "file"):
|
||
|
|
self.pbar = None
|
||
|
|
self.name = model_name
|
||
|
|
|
||
|
|
def __call__(self, block_num: int, block_size: int, total_size: int):
|
||
|
|
if not self.pbar:
|
||
|
|
self.pbar = tqdm(
|
||
|
|
desc=self.name,
|
||
|
|
initial=0,
|
||
|
|
unit="iB",
|
||
|
|
unit_scale=True,
|
||
|
|
unit_divisor=1000,
|
||
|
|
total=total_size,
|
||
|
|
)
|
||
|
|
self.pbar.update(block_size)
|
||
|
|
|
||
|
|
|
||
|
|
def download_with_progress_bar(name: str, url: str, dest_path: Path) -> bool:
|
||
|
|
"""Download a file from a URL to a destination path, with a progress bar.
|
||
|
|
If the file already exists, it will not be downloaded again.
|
||
|
|
|
||
|
|
Exceptions are not caught.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
name (str): Name of the file being downloaded.
|
||
|
|
url (str): URL to download the file from.
|
||
|
|
dest_path (Path): Destination path to save the file to.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
bool: True if the file was downloaded, False if it already existed.
|
||
|
|
"""
|
||
|
|
if dest_path.exists():
|
||
|
|
return False # already downloaded
|
||
|
|
|
||
|
|
InvokeAILogger.get_logger().info(f"Downloading {name}...")
|
||
|
|
|
||
|
|
dest_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
request.urlretrieve(url, dest_path, ProgressBar(name))
|
||
|
|
|
||
|
|
return True
|