| 54 | |
| 55 | |
| 56 | def cache_model_weights(hf_model_id: str, cache_dir: str): |
| 57 | from huggingface_hub import snapshot_download |
| 58 | |
| 59 | logging.info(f"Downloading model {hf_model_id}") |
| 60 | |
| 61 | if os.path.exists(cache_dir): |
| 62 | logging.info("Model already downloaded.") |
| 63 | |
| 64 | else: |
| 65 | os.makedirs(cache_dir, exist_ok=True) |
| 66 | |
| 67 | snapshot_download( |
| 68 | hf_model_id, local_dir=cache_dir, local_dir_use_symlinks=False |
| 69 | ) |
| 70 | logging.info("Committing model to volume.") |
| 71 | stub.volume.commit() |
| 72 | |
| 73 | logging.info("Model committed.") |
| 74 | |
| 75 | |
| 76 | def read_all_files(directory): |