Create upload-onnx.py

upload-onnx.py

@@ -0,0 +1,67 @@
+import os
+import shutil
+from huggingface_hub import HfApi, HfFolder, repo_type_and_id_from_hf_id, upload_file
+
+def main():
+    # Obtiene el nombre del modelo de la variable de entorno
+    model_name = os.getenv('MODELNAME')
+    if model_name is None:
+        print("MODELNAME no está definida en las variables de entorno.")
+        return
+
+    # Ruta a los archivos
+    onnx_file = f"{model_name}.onnx"
+    config_file = "config.json"
+    onnx_json_file = f"{model_name}.onnx.json"
+
+    # Verifica si el archivo .onnx existe
+    if os.path.isfile(onnx_file):
+        print(f"{model_name}.onnx encontrado.")
+    else:
+        print(f"{model_name}.onnx no se encuentra.")
+        return
+
+    # Renombra config.json a ${MODELNAME}.onnx.json si existe
+    if os.path.isfile(config_file):
+        shutil.move(config_file, onnx_json_file)
+        print(f"Renombrado config.json a {onnx_json_file}")
+    else:
+        print("config.json no se encuentra.")
+        return
+
+    # Configuración para la API de Hugging Face
+    token = HfFolder.get_token()  # Obtiene el token de acceso desde el archivo .huggingface/token
+    api = HfApi()
+
+    # Crea el repositorio en Hugging Face si no existe
+    repo_id = f"{HfFolder.get_username(token)}/{model_name}"
+    try:
+        api.create_repo(repo_id, token=token)
+        print(f"Repositorio {repo_id} creado en Hugging Face.")
+    except Exception as e:
+        print(f"El repositorio {repo_id} ya existe o se produjo un error: {e}")
+
+    # Subir archivos a Hugging Face
+    try:
+        # Subir el archivo .onnx
+        upload_file(
+            path_or_fileobj=onnx_file,
+            path_in_repo=f"{model_name}.onnx",
+            repo_id=repo_id,
+            token=token,
+        )
+        print(f"{model_name}.onnx subido a Hugging Face.")
+
+        # Subir el archivo .onnx.json
+        upload_file(
+            path_or_fileobj=onnx_json_file,
+            path_in_repo=f"{model_name}.onnx.json",
+            repo_id=repo_id,
+            token=token,
+        )
+        print(f"{model_name}.onnx.json subido a Hugging Face.")
+    except Exception as e:
+        print(f"Error al subir archivos a Hugging Face: {e}")
+
+if __name__ == '__main__':
+    main()