Check user login status

app.py

@@ -8,6 +8,7 @@ from io import BytesIO
 from PIL import Image
 import uuid
 import base64
+import json
 
 example_path = os.path.join(os.path.dirname(__file__), 'assets')
 clothing_list = os.listdir(os.path.join(example_path, "clothing"))
@@ -23,6 +24,8 @@ secret_key = os.getenv('secret_key')
 agent_version = os.getenv('agent_version')
 agent_name = os.getenv('agent_name')
 app_id = os.getenv('app_id')
+login_status_key = os.getenv('login_status_key')
+login_info_key = os.getenv('login_info_key')
 
 
 def parse_response(response, state='default'):
@@ -223,9 +226,48 @@ def image_to_base64(image):
     return f"data:image/png;base64,{img_base64}"
 
 
+def check_login_status(headers):
+    if not headers:
+        return False
+
+    try:
+        text = headers.get(login_status_key)
+        if not text or "." not in text:
+            return False
+
+        infos = text.split(".")
+        if len(infos) < 2:
+            return False
+
+        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)
+
+        data = datas.get(login_info_key)
+        if not data:
+            return False
+
+        user_id = data.get("_id")
+        user_name = data.get("user")
+        return bool(user_id and user_name)
+
+    except Exception as e:
+        print(f"An error occurred: {repr(e)}")
+        return False
+
+
 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:
+        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):
         m = "Please upload the main image before generating."
         return gr.Warning(m), did