Fix login

README.md

@@ -9,6 +9,7 @@ app_file: app.py
 pinned: false
 license: other
 short_description: WeShopAI Bad Hand Fixer. Fix hand areas in photos easily.
+hf_oauth: true
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -9,6 +9,7 @@ from PIL import Image
 import uuid
 import base64
 import json
+from huggingface_hub import whoami
 
 example_path = os.path.join(os.path.dirname(__file__), 'assets')
 clothing_list = os.listdir(os.path.join(example_path, "clothing"))
@@ -240,62 +241,27 @@ def image_to_base64(image):
     return f"data:image/png;base64,{img_base64}"
 
 
-def check_login_status(headers):
-    if not headers:
-        return None, None
-
-    try:
-        login_status = 1
-        text = headers.get(login_status_key)
-        if not text or "." not in text:
-            text = headers.get(login_status_key2)
-            login_status = 2
-            if not text or "." not in text:
-                return None, None
-
-        infos = text.split(".")
-        if len(infos) < 2:
-            return None, None
-
-        info = infos[1]
-        cover = len(info) % 4
-        if cover != 0:
-            info += "=" * (4 - cover)
-
-        decoded_bytes = base64.b64decode(info)
-        decoded_str = decoded_bytes.decode('utf-8')
-        datas = json.loads(decoded_str)
-        if login_status == 1:
-            data = datas.get(login_info_key)
-            if not data:
-                return None, None
-            user_id = data.get("_id")
-            user_name = data.get("user")
-            return user_id, user_name
-        elif login_status == 2:
-            user_id = datas.get("uuid")
-            user_name = datas.get("user")
-            return user_id, user_name
-        else:
-            return None, None
-
-    except Exception as e:
-        print(f"An error occurred: {repr(e)}")
-        return None, None
-
-
-def generate_image(edit_image_infos, did, request: gr.Request):
+def generate_image(edit_image_infos, did, request: gr.Request, oauth_token: gr.OAuthToken | None):
     if not did:
         did = str(uuid.uuid4())
-    user_id, user_name = check_login_status(request.request.headers)
-    if not user_id:
-        user_id = request.session_hash
-    if not user_name:
-        user_name = request.session_hash
+    if not oauth_token:
+        m = "Please click the authorization button below the Generate button to authorize."
+        return gr.Warning(m), did
+    user_id = None
+    user_name = None
+    try:
+        oauth_infos = whoami(oauth_token.token)
+        print(f"oauth_infos={oauth_infos}")
+        if oauth_infos and isinstance(oauth_infos, dict):
+            user_id = oauth_infos.get('id')
+            user_name = oauth_infos.get('name')
+    except Exception as e:
+        print(f"error: {repr(e)}")
+        m = "There was an error processing the authorization information. Please refresh the page and try again."
+        return gr.Warning(m), did
     if not user_id or not user_name:
-        m = "Please log in to your Hugging Face account to use the features of this application."
+        m = "Missing user login information. Please check and try refreshing the page."
         return gr.Warning(m), did
-    user_id = f"{login_hash_key}{user_id}"
     if edit_image_infos is None or not isinstance(edit_image_infos, dict):
         m = "Please upload the main image before generating."
         return gr.Warning(m), did
@@ -544,6 +510,8 @@ with gr.Blocks(css=css) as WeShop:
                     outputs=[output, current_did],
                     concurrency_limit=None
                 )
+            with gr.Row():
+                gr.LoginButton()
             with gr.Column():
                 main_image_middleware = gr.Image(
                     image_mode='RGBA',