Spaces:
Runtime error
Runtime error
import os | |
from typing import Optional | |
import requests | |
from tqdm.auto import tqdm | |
MODEL_PATHS = { | |
"clevr": "https://www.dropbox.com/s/bqpc3ymstz9m05z/clevr_model.pt", | |
"celebahq": "https://www.dropbox.com/s/687wuamoud4cs9x/celeb_model.pt" | |
} | |
DATA_PATHS = { | |
"clevr": "", | |
"clevr_toy": "" | |
} | |
def download_model( | |
dataset: str, | |
cache_dir: Optional[str] = None, | |
chunk_size: int = 4096, | |
) -> str: | |
if dataset not in MODEL_PATHS: | |
raise ValueError( | |
f"Unknown dataset name {dataset}. Known names are: {MODEL_PATHS.keys()}." | |
) | |
if cache_dir is None: | |
cache_dir = './' | |
url = MODEL_PATHS[dataset] | |
os.makedirs(cache_dir, exist_ok=True) | |
local_path = os.path.join(cache_dir, url.split("/")[-1]) | |
if os.path.exists(local_path.replace('?dl=0', '')): | |
return local_path.replace('?dl=0', '') | |
headers = {'user-agent': 'Wget/1.16 (linux-gnu)'} | |
r = requests.get(url, stream=True, headers=headers) | |
size = int(r.headers.get("content-length", "0")) | |
with open(local_path, 'wb') as f: | |
pbar = tqdm(total=size, unit="iB", unit_scale=True) | |
for chunk in r.iter_content(chunk_size=chunk_size): | |
if chunk: | |
pbar.update(len(chunk)) | |
f.write(chunk) | |
os.rename(local_path, local_path.replace('?dl=0', '')) | |
return local_path.replace('?dl=0', '') |