Update main.py

main.py

@@ -4,14 +4,15 @@ import uuid
 import json
 import os
 import requests
-import httpx  # New import for async requests
-import asyncio # New import for the background task
+import httpx
+import asyncio
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from PIL import Image
 import logging
 from io import BytesIO
 import time
+from datetime import datetime, timedelta # New imports for date manipulation
 
 # --- Configuration ---
 # Set up logging for debugging
@@ -25,8 +26,8 @@ PROXY_REFRESH_INTERVAL = 5 * 60  # 5 minutes in seconds
 # --- FastAPI App Initialization ---
 app = FastAPI(
     title="Image Generation API",
-    description="An API to generate images with periodic proxy list refreshing.",
-    version="1.3.0" # Version bumped for the new feature
+    description="An API to generate images with periodic proxy refreshing and randomized headers.",
+    version="1.4.0" # Version bumped for the new feature
 )
 
 # ----------- Proxy Management -----------
@@ -54,17 +55,14 @@ async def fetch_proxies():
             logger.warning("Proxy list was fetched but is empty.")
             PROXY_LIST = []
         else:
-            # This is an atomic update, so it's safe in this asyncio context
             PROXY_LIST = proxies
             logger.info(f"Successfully refreshed proxy list. Loaded {len(PROXY_LIST)} proxies.")
 
     except httpx.RequestError as e:
         logger.error(f"Failed to fetch or refresh proxy list: {e}")
-        # We don't clear the list on failure, to keep using old proxies if the source is temporarily down
     except Exception as e:
         logger.error(f"An unexpected error occurred during proxy fetch: {e}")
 
-
 async def schedule_proxy_refresh():
     """
     A background task that runs forever, refreshing the proxy list periodically.
@@ -85,17 +83,12 @@ async def startup_event():
     asyncio.create_task(schedule_proxy_refresh())
 
 def get_random_proxy():
-    """
-    Selects a random proxy from the global list and formats it for the requests library.
-    Returns None if the proxy list is empty.
-    """
+    """Selects a random proxy from the global list."""
     if not PROXY_LIST:
         return None
-    
     proxy = random.choice(PROXY_LIST)
     return {"http": proxy, "https": proxy}
 
-
 # ----------- Models -----------
 class GenerationRequest(BaseModel):
     prompt: str
@@ -114,6 +107,17 @@ def generate_session_hash():
 def generate_zerogpu_uuid():
     return f"n{random.randint(1000,9999)}_{uuid.uuid4().hex[:12]}_creitl"
 
+def generate_random_future_date_header():
+    """
+    Generates a random date string for an HTTP header, between now and 24 hours in the future.
+    """
+    # Add a random number of seconds (up to 24 hours) to the current time
+    random_seconds = random.randint(0, 24 * 3600)
+    future_time = datetime.utcnow() + timedelta(seconds=random_seconds)
+    # Format according to RFC 7231 for HTTP 'Date' header
+    return future_time.strftime('%a, %d %b %Y %H:%M:%S GMT')
+
+
 # ----------- Helper Function to Upload to SnapZion -----------
 def upload_to_snapzion(file_content: BytesIO, file_name: str):
     token = os.getenv("SNAPZION_API_TOKEN")
@@ -141,10 +145,16 @@ async def generate_image(req: GenerationRequest):
     session_hash = req.session_hash if req.session_hash else generate_session_hash()
     uuid_token = generate_zerogpu_uuid()
 
+    # --- HEADER CUSTOMIZATION ---
+    # Generate a new random Date header for each request
+    random_date_header = generate_random_future_date_header()
+    logger.info(f"Using randomized Date header for this request: {random_date_header}")
+
     headers_post = {
         "accept": "*/*",
         "accept-language": "en-US,en;q=0.9",
         "content-type": "application/json",
+        "Date": random_date_header, # <-- Here is the new randomized header
         "origin": "https://heartsync-nsfw-uncensored.hf.space",
         "referer": "https://heartsync-nsfw-uncensored.hf.space/?not-for-all-audiences=true&__theme=system",
         "sec-ch-ua": '"Not/A)Brand";v="8", "Chromium";v="126", "Google Chrome";v="126"',
@@ -166,7 +176,7 @@ async def generate_image(req: GenerationRequest):
     }
 
     try:
-        # Step 1: Join Queue (with proxy)
+        # Step 1: Join Queue
         join_proxy = get_random_proxy()
         proxy_ip_log = join_proxy['http'].split('@')[-1] if join_proxy else 'None'
         logger.info(f"Step 1: Joining queue with proxy: {proxy_ip_log}")
@@ -178,13 +188,13 @@ async def generate_image(req: GenerationRequest):
         join_response.raise_for_status()
         logger.info("Successfully joined queue. Now listening for SSE data.")
 
-        # Step 2: Listen to queue via SSE (with another proxy)
+        # Step 2: Listen to queue via SSE
         sse_proxy = get_random_proxy()
         proxy_ip_log = sse_proxy['http'].split('@')[-1] if sse_proxy else 'None'
         logger.info(f"Step 2: Listening to SSE stream with proxy: {proxy_ip_log}")
         
         sse_url = f"https://heartsync-nsfw-uncensored.hf.space/gradio_api/queue/data?session_hash={session_hash}"
-        sse_headers = {"accept": "text/event-stream", "user-agent": headers_post["user-agent"]}
+        sse_headers = {"accept": "text/event-stream", "user-agent": headers_post["user-agent"], "Date": random_date_header}
 
         with requests.get(sse_url, headers=sse_headers, stream=True, proxies=sse_proxy, timeout=300) as sse_response:
             sse_response.raise_for_status()
@@ -194,9 +204,7 @@ async def generate_image(req: GenerationRequest):
                 if decoded_line.startswith("data:"):
                     try:
                         data = json.loads(decoded_line[5:].strip())
-                        msg = data.get("msg")
-                        
-                        if msg == "process_completed":
+                        if data.get("msg") == "process_completed":
                             logger.info("Process completed. Extracting image data.")
                             output_list = data.get("output", {}).get("data", [])
                             if not output_list or not isinstance(output_list, list) or len(output_list) == 0:
@@ -208,7 +216,7 @@ async def generate_image(req: GenerationRequest):
                             if not image_url:
                                 raise HTTPException(status_code=500, detail="Image data received, but the URL is missing.")
 
-                            # Step 3: Download image (with a third proxy)
+                            # Step 3: Download image
                             download_proxy = get_random_proxy()
                             proxy_ip_log = download_proxy['http'].split('@')[-1] if download_proxy else 'None'
                             logger.info(f"Step 3: Downloading image with proxy: {proxy_ip_log}")
@@ -228,7 +236,6 @@ async def generate_image(req: GenerationRequest):
                             logger.info(f"Image successfully uploaded to SnapZion: {snapzion_url}")
                             
                             return {"success": True, "image_url": snapzion_url}
-
                     except (json.JSONDecodeError, IndexError): continue
 
         raise HTTPException(status_code=504, detail="Stream closed before generation could complete.")