Changing upload.py to accept model path as argument

models/ESD-X/esd-kellymckernan_from_kellymckernan-xattn_0-epochs_200.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:68e6e320371cf258ab68976ec33ef9d3cc6def9d0b0a8953266235efa2ce6fd9
-size 175884838

upload.py

@@ -1,7 +1,9 @@
-from huggingface_hub import HfApi
 import os
-from dotenv import load_dotenv
 import logging
+from huggingface_hub import HfApi
+from dotenv import load_dotenv
+import argparse
+
 logging.basicConfig(level=logging.DEBUG)
 
 # Load environment variables from .env file
@@ -9,29 +11,28 @@ load_dotenv()
 
 # Access Hugging Face API token
 hf_access_token = os.getenv("HUGGINGFACE_API_KEY")
-print(hf_access_token)  # Print to confirm loading, remove in production
-
-# Base path to your models directory
-MODEL_PATH = "/Users/klu/ErasingDiffusionModels-1/models"
-repo_id = "ErasureResearch/ErasingDiffusionModels"  # Target Hugging Face repository
 
+# Set up Hugging Face API
 api = HfApi()
 
+# Repository ID to upload to
+repo_id = "ErasureResearch/ErasingDiffusionModels"  # Target Hugging Face repository
+
 # Ensure the repository exists
 api.create_repo(repo_id, exist_ok=True, token=hf_access_token)
 
 # List existing files in the repository to avoid overwrites
 existing_files = api.list_repo_files(repo_id=repo_id, repo_type="model")
 
-def upload_model_files(subfolder_path):
+def upload_model_files(model_path, subfolder_path):
     for root, _, files in os.walk(subfolder_path):
         for file in files:
             file_path = os.path.join(root, file)
-            path_in_repo = os.path.relpath(file_path, MODEL_PATH)
+            path_in_repo = os.path.relpath(file_path, model_path)
             
             # Check if the file already exists in the repository
             if path_in_repo not in existing_files:
-                print(f"Uploading {file_path} to {repo_id}...")
+                logging.info(f"Uploading {file_path} to {repo_id}...")
                 api.upload_file(
                     path_or_fileobj=file_path,
                     path_in_repo=path_in_repo,
@@ -39,15 +40,24 @@ def upload_model_files(subfolder_path):
                     repo_type="model",
                     token=hf_access_token
                 )
-                print(f"Successfully uploaded {file_path} to {repo_id}.")
+                logging.info(f"Successfully uploaded {file_path} to {repo_id}.")
             else:
-                print(f"File {path_in_repo} already exists in {repo_id}, skipping upload.")
+                logging.info(f"File {path_in_repo} already exists in {repo_id}, skipping upload.")
 
 def main():
-    for subfolder in os.listdir(MODEL_PATH):
-        subfolder_path = os.path.join(MODEL_PATH, subfolder)
+    # Argument parser for model path
+    parser = argparse.ArgumentParser(description="Upload model files to Hugging Face Hub")
+    parser.add_argument("model_path", type=str, help="Path to the models directory")
+    args = parser.parse_args()
+    
+    # Base path to your models directory
+    model_path = args.model_path
+
+    # Process each subfolder in the specified model path
+    for subfolder in os.listdir(model_path):
+        subfolder_path = os.path.join(model_path, subfolder)
         if os.path.isdir(subfolder_path):
-            upload_model_files(subfolder_path)
+            upload_model_files(model_path, subfolder_path)
 
 if __name__ == "__main__":
     main()