Update app.py

app.py

@@ -12,6 +12,7 @@ from io import BytesIO
 from googlesearch import search
 import logging
 import queue
+from huggingface_hub import HfApi
 
 # Create a logging filter to suppress socket warnings
 class SocketWarningFilter(logging.Filter):
@@ -126,33 +127,36 @@ def search_images(query, num_images=5):
 
 HF_TOKEN = os.getenv("HF_TOKEN")  # Make sure you set the HF_TOKEN in your environment
 
-def restart_space(space_id):
-    """Restart a Hugging Face Space using the API."""
-    api_url = f"https://api.huggingface.co/spaces/{space_id}/restart"  # Corrected API endpoint for restarting Spaces
-    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
-    try:
-        response = requests.post(api_url, headers=headers)
-        if response.status_code == 200:
-            logger.info(f"Successfully restarted Space: {space_id}")
-            return jsonify({"success": True, "message": f"Successfully restarted Space: {space_id}"}), 200
-        else:
-            logger.error(f"Failed to restart Space: {space_id} - {response.status_code} - {response.text}")
-            return jsonify({"success": False, "message": f"Failed to restart Space: {space_id} - {response.status_code} - {response.text}"}), 500
-    except Exception as e:
-        logger.error(f"Error restarting space: {str(e)}")
-        return jsonify({"success": False, "message": f"Error: {str(e)}"}), 500
-
 @app.route('/restart_space', methods=['POST'])
 def api_restart_space():
     """API route to restart a Hugging Face Space."""
     space_id = 'Pamudu13/web-scraper'  
-    
-    # Expecting the space_id in the JSON body
+    factory_reboot = request.json.get('factory_reboot', False)  # Optional: Set to True if you want a factory reboot
 
     if not space_id:
         return jsonify({'error': 'space_id parameter is required'}), 400
 
-    return restart_space(space_id)
+    try:
+        hfapi = HfApi()
+
+        # Call the restart_space method
+        res = hfapi.restart_space(
+            space_id,
+            token=HF_TOKEN,
+            factory_reboot=factory_reboot
+        )
+
+        return jsonify({
+            'success': True,
+            'message': f"Successfully restarted Space: {space_id}",
+            'response': res
+        }), 200
+
+    except Exception as e:
+        return jsonify({
+            'success': False,
+            'message': f"Error: {str(e)}"
+        }), 500    
 
 @app.route('/search_images', methods=['GET'])
 def api_search_images():