decomp-diffusion / download.py
jsu27's picture
celeb combination demo
90bc5e7
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', '')