authentication

app.py

@@ -54,10 +54,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)
@@ -69,7 +71,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,
@@ -82,6 +84,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,
@@ -94,7 +98,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,
@@ -107,6 +111,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,
@@ -118,7 +124,7 @@ def create_task(image_url, did, ip):
     return data, msg
 
 
-def save_task(image_url, show_image, task_id, text_description, did, ip):
+def save_task(image_url, show_image, task_id, text_description, did, ip, user_id, user_name):
     timestamp = str(int(time.time()))
     signature = generate_signature(
         key=secret_key,
@@ -131,6 +137,8 @@ def save_task(image_url, show_image, task_id, text_description, 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,
@@ -145,7 +153,7 @@ def save_task(image_url, show_image, task_id, text_description, 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,
@@ -158,6 +166,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,
@@ -170,7 +180,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
     timestamp = str(int(time.time()))
@@ -192,6 +202,8 @@ 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_url, files=files, headers=headers)
         data, msg = parse_response(response)
@@ -206,16 +218,16 @@ def load_description(file_path):
 
 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
@@ -228,22 +240,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(main_image, text_description, 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 main_image is None:
@@ -259,6 +271,8 @@ def generate_image(main_image, text_description, did, request: gr.Request):
         image=main_image,
         did=did,
         ip=client_ip,
+        user_id=user_id,
+        user_name=user_name
     )
     if not upload_image_data:
         return gr.Warning(upload_image_msg), did
@@ -270,6 +284,8 @@ def generate_image(main_image, text_description, did, request: gr.Request):
         image_url=image_url,
         did=did,
         ip=client_ip,
+        user_id=user_id,
+        user_name=user_name
     )
     if not create_task_data:
         return gr.Warning(create_task_msg), did
@@ -285,6 +301,8 @@ def generate_image(main_image, text_description, did, request: gr.Request):
         text_description=text_description,
         did=did,
         ip=client_ip,
+        user_id=user_id,
+        user_name=user_name
     )
     if not save_task_data:
         return gr.Warning(save_task_msg), did
@@ -296,6 +314,8 @@ def generate_image(main_image, text_description, did, request: gr.Request):
         task_id=save_task_id,
         did=did,
         ip=client_ip,
+        user_id=user_id,
+        user_name=user_name
     )
     if not start_task_data:
         return gr.Warning(start_task_msg), did
@@ -311,6 +331,8 @@ def generate_image(main_image, text_description, did, request: gr.Request):
             execution_id=execution_id,
             did=did,
             ip=client_ip,
+            user_id=user_id,
+            user_name=user_name
         )
         print("Query task data:", query_task_data)
         print("Query task msg:", query_task_msg)
@@ -329,7 +351,7 @@ def generate_image(main_image, text_description, 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.'