sapiens-demo / download_checkpoints.py
joselobenitezg's picture
wip
94f04b7
import os
import json
import requests
from tqdm import tqdm
from config import SAPIENS_LITE_MODELS
def download_file(url, filename):
response = requests.get(url, stream=True)
total_size = int(response.headers.get('content-length', 0))
with open(filename, 'wb') as file, tqdm(
desc=filename,
total=total_size,
unit='iB',
unit_scale=True,
unit_divisor=1024,
) as progress_bar:
for data in response.iter_content(chunk_size=1024):
size = file.write(data)
progress_bar.update(size)
def main():
# Load the JSON file with model URLs
model_urls = SAPIENS_LITE_MODELS
for task, models in model_urls.items():
checkpoints_dir = os.path.join('checkpoints', task)
os.makedirs(checkpoints_dir, exist_ok=True)
for model_name, url in models.items():
model_filename = f"{model_name}_torchscript.pt2"
model_path = os.path.join(checkpoints_dir, model_filename)
if not os.path.exists(model_path):
print(f"Downloading {task} {model_name} model...")
download_file(url, model_path)
print(f"{task} {model_name} model downloaded successfully.")
else:
print(f"{task} {model_name} model already exists. Skipping download.")
if __name__ == "__main__":
main()