Update app.py

app.py

@@ -15,7 +15,7 @@ except:
 
 
 
-from flask import Flask, render_template, Response, request, jsonify, send_file
+from flask import Flask, render_template, Response, request, jsonify, send_file, abort
 import cv2
 import numpy as np
 import base64
@@ -26,6 +26,25 @@ import shutil
 
 app = Flask(__name__)
 
+# Get the access key from Hugging Face environment variable
+ACCESS_KEY = os.environ.get('KEY', 'default_access_key')
+
+def check_auth():
+    """Check if the request is authenticated with the correct key"""
+    provided_key = request.args.get('key')
+    if provided_key != ACCESS_KEY:
+        return False
+    return True
+
+def require_auth(f):
+    """Decorator to require authentication for routes"""
+    def decorated_function(*args, **kwargs):
+        if not check_auth():
+            abort(401)  # Unauthorized
+        return f(*args, **kwargs)
+    decorated_function.__name__ = f.__name__
+    return decorated_function
+
 # Initialize the face detection and swapping models
 app_model = FaceAnalysis(name='buffalo_l')
 app_model.prepare(ctx_id=0, det_size=(640, 640))
@@ -50,11 +69,13 @@ default_face = faces_in_default_img[0]
 
 
 @app.route('/')
+@require_auth
 def index():
     return render_template('index.html')
 
 
 @app.route('/process_frame', methods=['POST'])
+@require_auth
 def process_frame():
     try:
         data = request.json['image']
@@ -83,11 +104,13 @@ def process_frame():
 
 
 @app.route('/get_current_face')
+@require_auth
 def get_current_face():
     return send_file(default_img_path, mimetype='image/jpeg')
 
 
 @app.route('/upload_face', methods=['POST'])
+@require_auth
 def upload_face():
     if 'face' not in request.files:
         return jsonify({'success': False, 'message': 'No file part'})
@@ -122,6 +145,7 @@ def upload_face():
 
 
 @app.route('/reset_face', methods=['POST'])
+@require_auth
 def reset_face():
     backup_path = 'assets/face.jpg'
     if os.path.exists(backup_path):
@@ -138,5 +162,10 @@ def reset_face():
         return jsonify({'success': False, 'message': 'No backup found'})
 
 
+@app.errorhandler(401)
+def unauthorized(error):
+    return jsonify({'error': 'Unauthorized access. Valid key required.'}), 401
+
+
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)