authentication

app.py

@@ -61,10 +61,12 @@ def generate_signature(key, did, timestamp):
     return h.hexdigest()
 
 
-def url_to_image(url, ip):
+def url_to_image(url, ip, user_id, user_name):
     headers = {
         "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
-        'X-Forwarded-For': ip
+        'X-Forwarded-For': ip,
+        'X-HfUid': user_id,
+        'X-HfUserName': user_name,
     }
     try:
         response = requests.get(url, headers=headers, timeout=30)
@@ -76,7 +78,7 @@ def url_to_image(url, ip):
     return None
 
 
-def start_task(task_id, did, ip):
+def start_task(task_id, did, ip, user_id, user_name):
     timestamp = str(int(time.time()))
     signature = generate_signature(
         key=secret_key,
@@ -89,6 +91,8 @@ def start_task(task_id, did, ip):
         'X-Signature': signature,
         'X-Forwarded-For': ip,
         'X-AppId': app_id,
+        'X-HfUid': user_id,
+        'X-HfUserName': user_name,
     }
     data = {
         "agentVersion": agent_version,
@@ -101,7 +105,7 @@ def start_task(task_id, did, ip):
     return data, msg
 
 
-def create_task(image_url, did, ip):
+def create_task(image_url, did, ip, user_id, user_name):
     timestamp = str(int(time.time()))
     signature = generate_signature(
         key=secret_key,
@@ -114,6 +118,8 @@ def create_task(image_url, did, ip):
         'X-Signature': signature,
         'X-Forwarded-For': ip,
         'X-AppId': app_id,
+        'X-HfUid': user_id,
+        'X-HfUserName': user_name,
     }
     data = {
         "agentVersion": agent_version,
@@ -125,7 +131,7 @@ def create_task(image_url, did, ip):
     return data, msg
 
 
-def query_task(task_id, execution_id, did, ip):
+def query_task(task_id, execution_id, did, ip, user_id, user_name):
     timestamp = str(int(time.time()))
     signature = generate_signature(
         key=secret_key,
@@ -138,6 +144,8 @@ def query_task(task_id, execution_id, did, ip):
         'X-Signature': signature,
         'X-Forwarded-For': ip,
         'X-AppId': app_id,
+        'X-HfUid': user_id,
+        'X-HfUserName': user_name,
     }
     data = {
         "agentVersion": agent_version,
@@ -150,7 +158,7 @@ def query_task(task_id, execution_id, did, ip):
     return data, msg
 
 
-def upload_image(image, did, ip):
+def upload_image(image, did, ip, user_id, user_name):
     if image is None:
         return None
     image_format = image.format if image.format else "PNG"
@@ -171,13 +179,15 @@ def upload_image(image, did, ip):
             'X-Signature': signature,
             'X-Forwarded-For': ip,
             'X-AppId': app_id,
+            'X-HfUid': user_id,
+            'X-HfUserName': user_name,
         }
         response = requests.post(base_url + upload_image_url, files=files, headers=headers)
         data, msg = parse_response(response)
         return data, msg
 
 
-def mask_image_save(task_id, mask, did, ip):
+def mask_image_save(task_id, mask, did, ip, user_id, user_name):
     timestamp = str(int(time.time()))
     signature = generate_signature(
         key=secret_key,
@@ -190,6 +200,8 @@ def mask_image_save(task_id, mask, did, ip):
         'X-Signature': signature,
         'X-Forwarded-For': ip,
         'X-AppId': app_id,
+        'X-HfUid': user_id,
+        'X-HfUserName': user_name,
     }
     data = {
         "taskId": task_id,
@@ -228,16 +240,16 @@ def image_to_base64(image):
 
 def check_login_status(headers):
     if not headers:
-        return False
+        return None, None
 
     try:
         text = headers.get(login_status_key)
         if not text or "." not in text:
-            return False
+            return None, None
 
         infos = text.split(".")
         if len(infos) < 2:
-            return False
+            return None, None
 
         info = infos[1]
         cover = len(info) % 4
@@ -250,22 +262,22 @@ def check_login_status(headers):
 
         data = datas.get(login_info_key)
         if not data:
-            return False
+            return None, None
 
         user_id = data.get("_id")
         user_name = data.get("user")
-        return bool(user_id and user_name)
+        return user_id, user_name
 
     except Exception as e:
         print(f"An error occurred: {repr(e)}")
-        return False
+        return None, None
 
 
 def generate_image(edit_image_infos, did, request: gr.Request):
     if not did:
         did = str(uuid.uuid4())
-    login_status = check_login_status(request.request.headers)
-    if not login_status:
+    user_id, user_name = check_login_status(request.request.headers)
+    if not user_id or not user_name:
         m = "Please log in to your Hugging Face account to use the features of this application."
         return gr.Warning(m), did
     if edit_image_infos is None or not isinstance(edit_image_infos, dict):
@@ -296,7 +308,9 @@ def generate_image(edit_image_infos, did, request: gr.Request):
     upload_image_data, upload_image_msg = upload_image(
         image=main_image,
         did=did,
-        ip=client_ip
+        ip=client_ip,
+        user_id=user_id,
+        user_name=user_name
     )
     if not upload_image_data:
         return gr.Warning(upload_image_msg), did
@@ -307,7 +321,9 @@ def generate_image(edit_image_infos, did, request: gr.Request):
     create_task_data, create_task_msg = create_task(
         image_url=image_url,
         did=did,
-        ip=client_ip
+        ip=client_ip,
+        user_id=user_id,
+        user_name=user_name
     )
     if not create_task_data:
         return gr.Warning(create_task_msg), did
@@ -320,7 +336,9 @@ def generate_image(edit_image_infos, did, request: gr.Request):
         task_id=task_id,
         mask=image_to_base64(mask_image),
         did=did,
-        ip=client_ip
+        ip=client_ip,
+        user_id=user_id,
+        user_name=user_name
     )
     if not mask_image_save_data:
         return gr.Warning(mask_image_save_msg), did
@@ -328,7 +346,9 @@ def generate_image(edit_image_infos, did, request: gr.Request):
     start_task_data, start_task_msg = start_task(
         task_id=task_id,
         did=did,
-        ip=client_ip
+        ip=client_ip,
+        user_id=user_id,
+        user_name=user_name
     )
     if not start_task_data:
         return gr.Warning(start_task_msg), did
@@ -343,7 +363,9 @@ def generate_image(edit_image_infos, did, request: gr.Request):
             task_id=task_id,
             execution_id=execution_id,
             did=did,
-            ip=client_ip
+            ip=client_ip,
+            user_id=user_id,
+            user_name=user_name
         )
         if not query_task_data:
             return gr.Warning(query_task_msg), did
@@ -360,7 +382,7 @@ def generate_image(edit_image_infos, did, request: gr.Request):
         elif status == "Success" or status == "Blocked":
             img = results[0].get("image")
             if img and str(img).strip() != "":
-                return url_to_image(img, ip=client_ip), did
+                return url_to_image(img, ip=client_ip, user_id=user_id, user_name=user_name), did
         end_time = int(time.time())
         if end_time - start_time > 3600:
             m = 'Query task timeout.'

assets/title.html

@@ -93,7 +93,7 @@
             Explore our other applications on Hugging Face Spaces:
             <span class="nav-links">
                 <a href="https://huggingface.co/spaces/WeShopAI/WeShopAI-Virtual-Try-On">WeShopAI Virtual Try-On</a>
-                <a href="https://huggingface.co/spaces/WeShopAI/WeShopAI-Pose-Generator">WeShopAI Pose Generator</a>
+                <a href="https://huggingface.co/spaces/WeShopAI/WeShopAI-Fashion-Model-Pose-Change">WeShopAI Fashion Model Pose Change</a>
             </span>
         </div>
         <br>