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')
 human_list = os.listdir(os.path.join(example_path, "human"))
@@ -25,6 +23,9 @@ secret_key = os.getenv('secret_key')
 app_id = os.getenv('app_id')
 agent_version = os.getenv('agent_version')
 agent_name = os.getenv('agent_name')
+login_status_key = os.getenv('login_status_key')
+login_info_key = os.getenv('login_info_key')
+
 
 def parse_response(response):
     data = {}
@@ -203,9 +204,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, 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:
+        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:
         m = "Please upload both the main image and the background reference image before generating."
         return gr.Warning(m), did
@@ -214,7 +254,7 @@ def generate_image(main_image, text_description, did, request: gr.Request):
     x_forwarded_for = request.headers.get('x-forwarded-for')
     if x_forwarded_for:
         client_ip = x_forwarded_for
-   
+
     upload_image_data, upload_image_msg = upload_image(
         image=main_image,
         did=did,
@@ -384,7 +424,7 @@ with gr.Blocks(css=css) as WeShop:
                     "Low angle upward shot",
                 ],
                 outputs=text_description_input
-           )
+            )
         with gr.Column():
             with gr.Row():
                 with gr.Column():
@@ -406,22 +446,22 @@ with gr.Blocks(css=css) as WeShop:
                 show_case = gr.Examples(
                     examples=[
                         [
-                            "assets/human/07.png", 
-                            "Walking angle", 
+                            "assets/human/07.png",
+                            "Walking angle",
                             "assets/examples/result_07_01.jpeg",
                         ],
                         [
-                            "assets/human/01.png", 
+                            "assets/human/01.png",
                             "Following camera angle",
                             "assets/examples/result_01_01.jpeg"
                         ],
                         [
-                            "assets/human/06.png", 
+                            "assets/human/06.png",
                             "Change pose",
                             "assets/examples/result_06_01.jpeg"
                         ],
                         [
-                            "assets/human/04.png", 
+                            "assets/human/04.png",
                             "Keep the background unchanged, only change the person's pose",
                             "assets/examples/result_04_01.jpeg"
                         ],
@@ -434,4 +474,4 @@ with gr.Blocks(css=css) as WeShop:
         inputs=[main_image_input],
         outputs=main_image_input
     )
-WeShop.queue(api_open=False).launch(show_api=False)
+WeShop.queue(api_open=False).launch(show_api=False)