Update app.py

app.py

@@ -14,9 +14,22 @@ def image_to_base64(image):
 def base64_to_image(base64_str):
     return Image.open(BytesIO(base64.b64decode(base64_str + '=' * (-len(base64_str) % 4))))
 
-def search_face(file):
+def search_face(file, token):
     global count    
-    url = os.environ.get("SERVER_URL")
+    free_url = os.environ.get("SERVER_URL_FREE")
+    premium_url = os.environ.get("SERVER_URL_PREMIUM")
+    token_server_url = os.environ.get("TOKEN_URL")
+
+    url = free_url
+    if token:
+        try:
+            r = requests.post(url=token_server_url + '/verify-token', json={"token": token})
+            if r.json().get('status') == 'success':
+                url = premium_url
+            else:
+                raise gr.Error("Invalid token! For free search, use empty string for token")
+        except:
+            raise gr.Error("Invalid token!")
 
     try:
         image = Image.open(file)
@@ -51,15 +64,35 @@ def search_face(file):
     try:        
         res = r.json().get('img_array')
         out_array = []
+        suffix = "*********"
+        if url == premium_url:
+            suffix = ""
         for item in res:
-            out_array.append((base64_to_image(item["image"]), item["url"] + "*********"))
-        count += 1            
+            out_array.append((base64_to_image(item["image"]), item["url"] + suffix))
+        count += 1
+
+        if url == premium_url:
+            try:
+                r = requests.post(url=token_server_url + '/activate-token', json={"token": token})
+                #if r.json().get('status') == 'success':
+                #    url = premium_url
+            except:
+                raise gr.Error("Invalid token!")
+
         return out_array, count
     except:
         raise gr.Error("Try again.")
 
+def search_face_free(file):
+    search_face(file, '')
 
-with gr.Blocks() as demo:
+custom_css = """
+caption.caption {
+    user-select: text;
+    cursor: text;
+}
+"""
+with gr.Blocks(css=custom_css) as demo:
     gr.Markdown(
         """
     # Search Your Face Online For Free
@@ -78,9 +111,9 @@ with gr.Blocks() as demo:
             output = gr.Gallery(label="Found Images", columns=[4], object_fit="contain", height="auto")
             countwg = gr.Number(label="Count")
     
-    gr.Examples(['examples/1.jpg', 'examples/2.jpg'], inputs=image, cache_examples=True, cache_mode='lazy', fn=search_face, outputs=[output, countwg])
+    gr.Examples(['examples/1.jpg', 'examples/2.jpg'], inputs=image, cache_examples=True, cache_mode='lazy', fn=search_face_free, outputs=[output, countwg])
 
-    search_face_button.click(search_face, inputs=image, outputs=[output, countwg], api_name=False)
+    search_face_button.click(search_face, inputs=[image, token], outputs=[output, countwg], api_name=False)
 
     gr.HTML('<a href="https://visitorbadge.io/status?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2FFaceOnLive%2FFace-Search-Online"><img src="https://api.visitorbadge.io/api/combined?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2FFaceOnLive%2FFace-Search-Online&labelColor=%23ff8a65&countColor=%2337d67a&style=flat&labelStyle=upper" /></a>')