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', '')