File size: 2,769 Bytes
f17ae24 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | import os
import argparse
from huggingface_hub import HfApi
def upload_to_hf():
api = HfApi()
# 1. Upload Franka Dataset
dataset_path = "/storage/ice-shared/ae8803che/hxue/data/dataset/franka"
dataset_repo = "t1an/world_model_dataset"
if os.path.exists(dataset_path):
print(f"Uploading Franka dataset to {dataset_repo}...")
try:
api.upload_folder(
folder_path=dataset_path,
repo_id=dataset_repo,
repo_type="dataset",
path_in_repo="franka",
# ignore_patterns=["*.pt"] # We want metadata.pt though
)
except Exception as e:
print(f"Error uploading dataset: {e}")
else:
print(f"Dataset path {dataset_path} not found.")
# 2. Upload Code and Latest Checkpoints to Model Repo
model_repo = "t1an/world_model"
# Create the repo if it doesn't exist
try:
api.create_repo(repo_id=model_repo, exist_ok=True)
except Exception as e:
print(f"Note on repo creation: {e}")
# Upload Code (wm folder)
print(f"Uploading code to {model_repo}...")
try:
api.upload_folder(
folder_path="/storage/ice-shared/ae8803che/hxue/data/world_model/wm",
repo_id=model_repo,
path_in_repo="wm"
)
except Exception as e:
print(f"Error uploading code: {e}")
# Upload only latest.pt from each checkpoint folder
checkpoint_dir = "/storage/ice-shared/ae8803che/hxue/data/world_model/checkpoints"
if os.path.exists(checkpoint_dir):
for run_name in os.listdir(checkpoint_dir):
run_path = os.path.join(checkpoint_dir, run_name)
if os.path.isdir(run_path):
latest_ckpt = os.path.join(run_path, "latest.pt")
if os.path.exists(latest_ckpt):
print(f"Uploading {run_name}/latest.pt to {model_repo}...")
try:
api.upload_file(
path_or_fileobj=latest_ckpt,
path_in_repo=f"checkpoints/{run_name}/latest.pt",
repo_id=model_repo
)
except Exception as e:
print(f"Error uploading {run_name}/latest.pt: {e}")
if __name__ == "__main__":
# Check if conversion is finished by checking for metadata.pt
metadata_path = "/storage/ice-shared/ae8803che/hxue/data/dataset/franka/metadata.pt"
if not os.path.exists(metadata_path):
print("Error: Franka dataset conversion not finished (metadata.pt missing).")
print("Please wait for the conversion to finish before running this script.")
else:
upload_to_hf()
|