Spaces:
Sleeping
Sleeping
| from huggingface_hub import hf_hub_download, hf_hub_url, get_hf_file_metadata | |
| from huggingface_hub.utils import disable_progress_bars | |
| from pathlib import Path | |
| from rich.progress import Progress | |
| from fire import Fire | |
| from typing import Union, List | |
| _EXPERTS = [ | |
| "10_model.pth", | |
| "Unified_learned_OCIM_RS200_6x+2x.pth", | |
| "dpt_hybrid-midas-501f0c75.pt", | |
| "icdar2015_hourglass88.pth", | |
| "model_final_e0c58e.pkl", | |
| "model_final_f07440.pkl", | |
| "scannet.pt", | |
| ] | |
| _MODELS = [ | |
| "vqa_prismer_base", | |
| "vqa_prismer_large", | |
| "vqa_prismerz_base", | |
| "vqa_prismerz_large", | |
| "caption_prismerz_base", | |
| "caption_prismerz_large", | |
| "caption_prismer_base", | |
| "caption_prismer_large", | |
| "pretrain_prismer_base", | |
| "pretrain_prismer_large", | |
| "pretrain_prismerz_base", | |
| "pretrain_prismerz_large", | |
| ] | |
| _REPO_ID = "lorenmt/prismer" | |
| def download_checkpoints( | |
| download_experts: bool = False, | |
| download_models: Union[bool, List] = False, | |
| hide_tqdm: bool = False, | |
| force_redownload: bool = False, | |
| ): | |
| if hide_tqdm: | |
| disable_progress_bars() | |
| # Convert to list and check for invalid names | |
| download_experts = _EXPERTS if download_experts else [] | |
| if download_models: | |
| # only download single model | |
| if isinstance(download_models, str): | |
| download_models = [download_models] | |
| assert all([m in _MODELS for m in download_models]), f"Invalid model name. Must be one of {_MODELS}" | |
| download_models = _MODELS if isinstance(download_models, bool) else download_models | |
| else: | |
| download_models = [] | |
| # Check if files already exist | |
| if not force_redownload: | |
| download_experts = [e for e in download_experts if not Path(f"./experts/expert_weights/{e}").exists()] | |
| download_models = [m for m in download_models if not Path(f"{m}/pytorch_model.bin").exists()] | |
| assert download_experts or download_models, "Nothing to download." | |
| with Progress() as progress: | |
| # Calculate total download size | |
| progress.print("[blue]Calculating download size...") | |
| total_size = 0 | |
| for expert in download_experts: | |
| url = hf_hub_url( | |
| filename=expert, | |
| repo_id=_REPO_ID, | |
| subfolder="expert_weights" | |
| ) | |
| total_size += get_hf_file_metadata(url).size | |
| for model in download_models: | |
| url = hf_hub_url( | |
| filename=f"pytorch_model.bin", | |
| repo_id=_REPO_ID, | |
| subfolder=model | |
| ) | |
| total_size += get_hf_file_metadata(url).size | |
| progress.print(f"[blue]Total download size: {total_size / 1e9:.2f} GB") | |
| # Download files | |
| total_files = len(download_experts) + len(download_models) | |
| total_task = progress.add_task(f"[green]Downloading files", total=total_files) | |
| if download_experts: | |
| expert_task = progress.add_task( | |
| f"[green]Downloading experts...", total=len(download_experts) | |
| ) | |
| out_folder = Path("experts/expert_weights") | |
| out_folder.mkdir(parents=True, exist_ok=True) | |
| for expert in download_experts: | |
| path = Path(hf_hub_download( | |
| filename=expert, | |
| repo_id=_REPO_ID, | |
| subfolder="expert_weights" | |
| )) | |
| path.resolve().rename(out_folder/path.name) | |
| path.unlink() | |
| progress.advance(expert_task) | |
| progress.advance(total_task) | |
| if download_models: | |
| model_task = progress.add_task( | |
| f"[green]Downloading models...", total=len(download_models) | |
| ) | |
| for model in download_models: | |
| path = Path(hf_hub_download( | |
| filename=f"pytorch_model.bin", | |
| repo_id=_REPO_ID, | |
| subfolder=model | |
| )) | |
| out_folder = Path("./logging")/model | |
| out_folder.mkdir(parents=True, exist_ok=True) | |
| path.resolve().rename(out_folder/"pytorch_model.bin") | |
| path.unlink() | |
| progress.advance(model_task) | |
| progress.advance(total_task) | |
| progress.print("[green]Done!") | |
| if __name__ == "__main__": | |
| Fire(download_checkpoints) | |