Check user login status

app.py

@@ -7,10 +7,8 @@ import requests
 from io import BytesIO
 from PIL import Image
 import uuid
-
-# from dotenv import load_dotenv
-
-# load_dotenv()
+import base64
+import json
 
 example_path = os.path.join(os.path.dirname(__file__), 'assets')
 clothing_list = os.listdir(os.path.join(example_path, "clothing"))
@@ -28,6 +26,9 @@ 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):
     data = {}
@@ -209,9 +210,48 @@ def load_description(file_path):
     return content
 
 
+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(main_image, background_image, 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 main_image is None or background_image is None:
         m = "Please upload both the main image and the background reference image before generating."
         return gr.Warning(m), did