Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
-import requests
-import base64
+from gradio_client import Client, GradioClientError
 from PIL import Image
+import base64
 import io
 import logging
 
@@ -10,62 +10,71 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(level
 logger = logging.getLogger(__name__)
 
 # ==============================================================================
-#  CLIENT-SIDE API CALL LOGIC
+#  CLIENT-SIDE API CALL LOGIC (Using the Official Gradio Client)
 # ==============================================================================
-def call_api(server_base_url: str, image: Image.Image) -> dict:
+
+def call_api(server_space_id: str, image: Image.Image) -> dict:
     """
-    Takes a PIL image, converts it to base64, sends it to the specified server API,
-    and returns the decrypted JSON data.
+    Uses the official gradio_client to call the secure decoder API.
     """
-    if not server_base_url or "hf.space" not in server_base_url:
-        raise gr.Error("Please provide a valid Hugging Face Space URL for the server API.")
+    # --- Input Validation ---
+    if not server_space_id or "/" not in server_space_id:
+        raise gr.Error("Please provide the Server Space ID in the format 'username/space-name'.")
     if image is None:
         raise gr.Error("Please upload an encrypted image.")
 
-    # --- CORRECTED API PATH ---
-    api_endpoint = f"{server_base_url.strip('/')}/run/keylock-auth-decoder"
-    logger.info(f"Client attempting to call API at: {api_endpoint}")
+    logger.info(f"Initializing client for server space: {server_space_id}")
 
     try:
+        # 1. Initialize the Gradio Client
+        # This points the client to your server space.
+        client = Client(src=server_space_id)
+
+        # 2. Convert the PIL Image to a base64 string
         with io.BytesIO() as buffer:
             image.save(buffer, format="PNG")
             image_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
 
-        payload = {"data": [image_base64]}
-        headers = {"Content-Type": "application/json"}
-        response = requests.post(api_endpoint, headers=headers, json=payload, timeout=30)
+        # 3. Call the specific API endpoint on the server
+        # This matches the documentation you provided exactly.
+        logger.info(f"Calling api_name='/keylock-auth-decoder' on {server_space_id}")
+        result = client.predict(
+            # The keyword argument MUST match the server function's parameter name.
+            image_base64_string=image_base64,
+            api_name="/keylock-auth-decoder"
+        )
         
-        response_json = response.json()
-        if response.status_code == 200:
-            if "data" in response_json:
-                return response_json["data"][0]
-            elif "error" in response_json:
-                raise gr.Error(f"API returned an error: {response_json['error']}")
-        else:
-            error_detail = response_json.get("error", "Unknown error.")
-            raise gr.Error(f"API Error (Status {response.status_code}): {error_detail}")
+        logger.info("Successfully received result from API.")
+        # The result from a successful predict call is the data itself.
+        return result
 
-    except requests.exceptions.RequestException as e:
-        logger.error(f"Network error calling API: {e}")
-        raise gr.Error(f"Could not connect to the API. Check the URL and server status. Error: {e}")
+    except GradioClientError as e:
+        # Handle errors from the gradio_client library (e.g., 404, server errors)
+        logger.error(f"Gradio Client Error: {e}")
+        raise gr.Error(f"API Call Failed: {e}. Check if the Server Space ID is correct and the server is running.")
     except Exception as e:
+        # Handle other unexpected errors
         logger.error(f"An unexpected client-side error occurred: {e}", exc_info=True)
         raise gr.Error(f"An unexpected error occurred: {e}")
 
 # ==============================================================================
 #  GRADIO INTERFACE
 # ==============================================================================
+
 with gr.Blocks(theme=gr.themes.Soft(), title="KeyLock API Client") as demo:
     gr.Markdown("# 🔑 KeyLock API Client")
-    gr.Markdown("This UI calls the **Secure Decoder API**. It sends an encrypted image to the server for decryption.")
+    gr.Markdown(
+        "This UI calls the **Secure Decoder API**. It sends an encrypted image to the server for decryption "
+        "using the official `gradio_client` library."
+    )
 
     with gr.Row():
         with gr.Column(scale=1):
-            gr.Markdown("### 1. Configure Server URL")
-            server_url_input = gr.Textbox(
-                label="Decoder API Server Base URL",
-                placeholder="https://your-name-keylock-auth.hf.space",
-                info="This is the direct URL of your server space, without any path."
+            gr.Markdown("### 1. Configure Server Space ID")
+            server_id_input = gr.Textbox(
+                label="Decoder Server Space ID",
+                value="broadfield-dev/KeyLock-Auth",
+                info="This is the Hugging Face ID of your server, e.g., 'username/space-name'."
             )
 
             gr.Markdown("### 2. Upload Encrypted Image")
@@ -78,7 +87,7 @@ with gr.Blocks(theme=gr.themes.Soft(), title="KeyLock API Client") as demo:
 
     submit_button.click(
         fn=call_api,
-        inputs=[server_url_input, image_input],
+        inputs=[server_id_input, image_input],
         outputs=[json_output]
     )