(name, *, path=None, url=None)
| 40 | pbar.update(len(chunk)) |
| 41 | |
| 42 | def auto_create(name, *, path=None, url=None): |
| 43 | if path is None: |
| 44 | path = os.getenv('SAT_HOME', '~/.sat_models') |
| 45 | path = os.path.expanduser(path) |
| 46 | model_path = os.path.join(path, name) |
| 47 | if url == 'local': |
| 48 | return model_path |
| 49 | os.makedirs(os.path.dirname(model_path), exist_ok=True) |
| 50 | lock = FileLock(model_path + '.lock', mode=0o777) |
| 51 | with lock: |
| 52 | if url is None: |
| 53 | url = MODEL_URLS[name] |
| 54 | if os.path.isdir(model_path) and not url.startswith('r2://'): |
| 55 | pass |
| 56 | elif os.path.isdir(model_path) and url.startswith('r2://') and url.endswith('.zip'): |
| 57 | pass |
| 58 | else: |
| 59 | print(f'Downloading models {url} into {path} ...') |
| 60 | try: |
| 61 | if url.startswith('r2://'): |
| 62 | download_s3(path, url[5:]) |
| 63 | else: |
| 64 | file_path = os.path.join(path, name + '.zip') |
| 65 | download_with_progress_bar(file_path, url) |
| 66 | except Exception as e: |
| 67 | print(f'Failed to download or check, if you already had the zip file, please unzip it manually as {model_path}!') |
| 68 | raise e |
| 69 | # unzip |
| 70 | if not os.path.isdir(model_path): |
| 71 | import zipfile |
| 72 | file_path = os.path.join(path, name + '.zip') |
| 73 | print(f'Unzipping {file_path}...') |
| 74 | f = zipfile.ZipFile(file_path, 'r') |
| 75 | f.extractall(path=path) |
| 76 | assert os.path.isdir(model_path), f'Unzip failed, or the first-level folder in zip is not {name}.' |
| 77 | return model_path # must return outside the `with lock` block |
| 78 | |
| 79 | SAT_ACCOUNT = 'c8a00746a80e06c4632028e37de24d6e' |
| 80 | SAT_ACCESS_KEY = 'eb4d69e273848089c7f9b9599cdcd983' |
no test coverage detected