update

app.py

@@ -52,9 +52,10 @@ def generate_signature(key, did, timestamp):
     return h.hexdigest()
 
 
-def url_to_image(url):
+def url_to_image(url, ip):
     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"
+        "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
     }
     try:
         response = requests.get(url, headers=headers, timeout=30)
@@ -66,11 +67,12 @@ def url_to_image(url):
     return None
 
 
-def start_task(task_id, timestamp, signature, did):
+def start_task(task_id, timestamp, signature, did, ip):
     headers = {
         'Did': did,
         'X-Timestamp': timestamp,
-        'X-Signature': signature
+        'X-Signature': signature,
+        'X-Forwarded-For': ip
     }
     data = {
         "agentVersion": agent_version,
@@ -83,11 +85,12 @@ def start_task(task_id, timestamp, signature, did):
     return data, msg
 
 
-def create_task(image_url, timestamp, signature, did):
+def create_task(image_url, timestamp, signature, did, ip):
     headers = {
         'Did': did,
         'X-Timestamp': timestamp,
-        'X-Signature': signature
+        'X-Signature': signature,
+        'X-Forwarded-For': ip
     }
     data = {
         "agentVersion": agent_version,
@@ -99,11 +102,12 @@ def create_task(image_url, timestamp, signature, did):
     return data, msg
 
 
-def save_task(image_url, timestamp, signature, show_image, task_id, location_id, did):
+def save_task(image_url, timestamp, signature, show_image, task_id, location_id, did, ip):
     headers = {
         'Did': did,
         'X-Timestamp': timestamp,
-        'X-Signature': signature
+        'X-Signature': signature,
+        'X-Forwarded-For': ip
     }
     data = {
         "agentVersion": agent_version,
@@ -118,11 +122,12 @@ def save_task(image_url, timestamp, signature, show_image, task_id, location_id,
     return data, msg
 
 
-def query_task(task_id, execution_id, timestamp, signature, did):
+def query_task(task_id, execution_id, timestamp, signature, did, ip):
     headers = {
         'Did': did,
         'X-Timestamp': timestamp,
-        'X-Signature': signature
+        'X-Signature': signature,
+        'X-Forwarded-For': ip
     }
     data = {
         "agentVersion": agent_version,
@@ -135,7 +140,7 @@ def query_task(task_id, execution_id, timestamp, signature, did):
     return data, msg
 
 
-def upload_image(image, signature, timestamp, upload_type, did):
+def upload_image(image, signature, timestamp, upload_type, did, ip):
     if image is None:
         return None
     if upload_type == 'image':
@@ -151,7 +156,8 @@ def upload_image(image, signature, timestamp, upload_type, did):
         headers = {
             'Did': did,
             'X-Timestamp': timestamp,
-            'X-Signature': signature
+            'X-Signature': signature,
+            'X-Forwarded-For': ip
         }
         response = requests.post(base_url + upload_url, files=files, headers=headers)
         data, msg = parse_response(response)
@@ -164,13 +170,15 @@ def load_description(file_path):
     return content
 
 
-def generate_image(main_image, background_image, did):
+def generate_image(main_image, background_image, did, request: gr.Request):
     if not did:
         did = str(uuid.uuid4())
     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
 
+    client_ip = request.client.host
+
     timestamp = str(int(time.time()))
     signature = generate_signature(
         key=secret_key,
@@ -182,7 +190,8 @@ def generate_image(main_image, background_image, did):
         signature=signature,
         timestamp=timestamp,
         upload_type='image',
-        did=did
+        did=did,
+        ip=client_ip
     )
     if not upload_image_data:
         return gr.Warning(upload_image_msg), did
@@ -194,7 +203,8 @@ def generate_image(main_image, background_image, did):
         image_url=image_url,
         timestamp=timestamp,
         signature=signature,
-        did=did
+        did=did,
+        ip=client_ip
     )
     if not create_task_data:
         return gr.Warning(create_task_msg), did
@@ -208,7 +218,8 @@ def generate_image(main_image, background_image, did):
         signature=signature,
         timestamp=timestamp,
         upload_type='background_image',
-        did=did
+        did=did,
+        ip=client_ip
     )
     if not upload_image_data:
         return gr.Warning(upload_image_msg), did
@@ -220,7 +231,8 @@ def generate_image(main_image, background_image, did):
         show_image=show_image,
         task_id=task_id,
         location_id=upload_image_data,
-        did=did
+        did=did,
+        ip=client_ip
     )
     if not save_task_data:
         return gr.Warning(save_task_msg), did
@@ -232,7 +244,8 @@ def generate_image(main_image, background_image, did):
         task_id=save_task_id,
         timestamp=timestamp,
         signature=signature,
-        did=did
+        did=did,
+        ip=client_ip
     )
     if not start_task_data:
         return gr.Warning(start_task_msg), did
@@ -248,7 +261,8 @@ def generate_image(main_image, background_image, did):
             execution_id=execution_id,
             timestamp=timestamp,
             signature=signature,
-            did=did
+            did=did,
+            ip=client_ip
         )
         if not query_task_data:
             return gr.Warning(query_task_msg), did
@@ -264,7 +278,7 @@ def generate_image(main_image, background_image, did):
         elif status == "Success" or status == "Blocked":
             img = results[0].get("image")
             if img and str(img).strip() != "":
-                return url_to_image(img), did
+                return url_to_image(img, ip=client_ip), did
         end_time = int(time.time())
         if end_time - start_time > 3600:
             m = 'Query task timeout.'