Update server.py

server.py

@@ -1,46 +1,39 @@
-from flask import Flask, send_from_directory, abort
 import os
+from flask import Flask, send_from_directory, abort
 
 app = Flask(__name__)
 
-# Definir la ruta donde se encuentran los archivos
-files_directory = "/home/app/"
+# Directorio donde se encuentran los archivos
+base_directory = "/home/app/"
 
-# Obtener el nombre del modelo de la variable de entorno
+# 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.")
-    # Si MODELNAME no está definida, asignar un valor por defecto
-    model_name = "default_model"
+    exit(1)
 
-# Ruta a los archivos
+# Ruta a los archivos basados en model_name
 onnx_file = f"{model_name}.onnx"
-onnx_json_file = f"{model_name}.onnx.json"
-
-# Verifica que los archivos existan en el directorio especificado
-def check_files_exist():
-    onnx_path = os.path.join(files_directory, onnx_file)
-    json_path = os.path.join(files_directory, onnx_json_file)
-    if not os.path.exists(onnx_path):
-        print(f"El archivo {onnx_file} no existe en {files_directory}")
-        abort(404, description=f"El archivo {onnx_file} no existe")
-    if not os.path.exists(json_path):
-        print(f"El archivo {onnx_json_file} no existe en {files_directory}")
-        abort(404, description=f"El archivo {onnx_json_file} no existe")
-
-
-# Ruta para visualizar el archivo .onnx
-@app.route('/download/onnx')
-def download_onnx():
-    check_files_exist()
-    return send_from_directory(files_directory, onnx_file)
-
-# Ruta para visualizar el archivo .onnx.json
-@app.route('/download/json')
-def download_onnx_json():
-    check_files_exist()
-    return send_from_directory(files_directory, onnx_json_file)
+json_file = f"{model_name}.onnx.json"
+
+@app.route('/onnx')
+def serve_onnx_file():
+    """Sirve el archivo .onnx."""
+    file_path = os.path.join(base_directory, onnx_file)
+    if os.path.exists(file_path):
+        return send_from_directory(base_directory, onnx_file)
+    else:
+        abort(404, description=f"El archivo {onnx_file} no se encuentra.")
+
+@app.route('/json')
+def serve_json_file():
+    """Sirve el archivo .onnx.json."""
+    file_path = os.path.join(base_directory, json_file)
+    if os.path.exists(file_path):
+        return send_from_directory(base_directory, json_file)
+    else:
+        abort(404, description=f"El archivo {json_file} no se encuentra.")
 
 if __name__ == '__main__':
-    # Iniciar el servidor Flask en el puerto 7860
+    # Ejecuta el servidor en el puerto 7860
     app.run(host='0.0.0.0', port=7860)