(url, ckpt_dir=None)
| 8 | |
| 9 | |
| 10 | def download(url, ckpt_dir=None): |
| 11 | name = url[url.rfind('/') + 1:] |
| 12 | if ckpt_dir is None: |
| 13 | ckpt_dir = tempfile.gettempdir() |
| 14 | ckpt_dir = os.path.join(ckpt_dir, 'art_fid') |
| 15 | ckpt_file = os.path.join(ckpt_dir, name) |
| 16 | if not os.path.exists(ckpt_file): |
| 17 | print(f'Downloading: \"{url[:url.rfind("?")]}\" to {ckpt_file}') |
| 18 | if not os.path.exists(ckpt_dir): |
| 19 | os.makedirs(ckpt_dir) |
| 20 | |
| 21 | response = requests.get(url, stream=True) |
| 22 | total_size_in_bytes = int(response.headers.get('content-length', 0)) |
| 23 | progress_bar = tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True) |
| 24 | |
| 25 | # first create temp file, in case the download fails |
| 26 | ckpt_file_temp = os.path.join(ckpt_dir, name + '.temp') |
| 27 | with open(ckpt_file_temp, 'wb') as file: |
| 28 | for data in response.iter_content(chunk_size=1024): |
| 29 | progress_bar.update(len(data)) |
| 30 | file.write(data) |
| 31 | progress_bar.close() |
| 32 | |
| 33 | if total_size_in_bytes != 0 and progress_bar.n != total_size_in_bytes: |
| 34 | print('An error occured while downloading, please try again.') |
| 35 | if os.path.exists(ckpt_file_temp): |
| 36 | os.remove(ckpt_file_temp) |
| 37 | else: |
| 38 | # if download was successful, rename the temp file |
| 39 | os.rename(ckpt_file_temp, ckpt_file) |
| 40 | return ckpt_file |
nothing calls this directly
no outgoing calls
no test coverage detected