Spaces:
Runtime error
Runtime error
File size: 1,379 Bytes
4c1d330 90bc5e7 4c1d330 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
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', '') |