upload-onnx.py

import os
import base64
import shutil
from huggingface_hub import HfApi, 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

    # Obtén el token desde la variable de entorno llamada 'TOKEN'
    token_base64 = os.getenv('TOKEN')
    if token_base64 is None:
        print("TOKEN no está definida en las variables de entorno.")
        return

    # Decodifica el token de base64
    token = base64.b64decode(token_base64).decode('utf-8')

    # Configuración para la API de Hugging Face
    api = HfApi()

    # Obtiene el nombre de usuario
    user_info = api.whoami(token=token)
    username = user_info['name']

    # Crea el repositorio en Hugging Face si no existe
    repo_id = f"{username}/{model_name}"
    try:
        api.create_repo(repo_id=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()