from __future__ import annotations import os import pathlib import shlex import subprocess import slugify from huggingface_hub import HfApi from constants import ( MODEL_LIBRARY_ORG_NAME, URL_TO_JOIN_MODEL_LIBRARY_ORG, UploadTarget, ) def join_model_library_org(hf_token: str) -> None: subprocess.run( shlex.split( f'curl -X POST -H "Authorization: Bearer {hf_token}" -H "Content-Type: application/json" {URL_TO_JOIN_MODEL_LIBRARY_ORG}' ) ) def upload( local_folder_path: str, target_repo_name: str, upload_to: str, private: bool = True, delete_existing_repo: bool = False, hf_token: str = "", ) -> str: hf_token = os.getenv("HF_TOKEN") or hf_token if not hf_token: raise ValueError api = HfApi(token=hf_token) if not local_folder_path: raise ValueError if not target_repo_name: target_repo_name = pathlib.Path(local_folder_path).name target_repo_name = slugify.slugify(target_repo_name) if upload_to == UploadTarget.PERSONAL_PROFILE.value: organization = api.whoami()["name"] elif upload_to == UploadTarget.MODEL_LIBRARY.value: organization = MODEL_LIBRARY_ORG_NAME join_model_library_org(hf_token) else: raise ValueError repo_id = f"{organization}/{target_repo_name}" if delete_existing_repo: try: api.delete_repo(repo_id, repo_type="model") except Exception: pass try: api.create_repo(repo_id, repo_type="model", private=private) api.upload_folder(repo_id=repo_id, folder_path=local_folder_path, path_in_repo=".", repo_type="model") url = f"https://huggingface.co/{repo_id}" message = f"Your model was successfully uploaded to {url}." except Exception as e: message = str(e) return message