Update gen.py

gen.py

@@ -42,7 +42,7 @@ PKEY3 = os.getenv("POLLINATIONS3_KEY", "")
 AIRFORCE_KEY = os.getenv("AIRFORCE")
 AIRFORCE_VIDEO_MODEL = "grok-imagine-video"
 AIRFORCE_API_URL = "https://api.airforce/v1/images/generations"
-
+WEBSOCKET_KEY = os.getenv("WEBSOCKET_KEY", "")
 valid_ratios = {"3:2", "2:3", "1:1", "", None}
 ratios = {"3:2", "2:3", "1:1"}
 valid_modes = {"normal", "fun", "", None}
@@ -123,7 +123,7 @@ async def generate_image(
 
     prompt = normalize_prompt_value(prompt, "prompt")
     enforce_prompt_size(prompt, MAX_MEDIA_PROMPT_CHARS, MAX_MEDIA_PROMPT_BYTES, "Image prompt")
-    if request.headers.get("x-payg-internal") != "true":
+    if request.headers.get("x-payg-internal") != WEBSOCKET_KEY:
         await check_image_rate_limit(request, authorization, x_client_id)
 
     chosen_model = "zimage"
@@ -191,7 +191,7 @@ async def gensfx(
 
     prompt = normalize_prompt_value(prompt, "prompt")
     enforce_prompt_size(prompt, MAX_MEDIA_PROMPT_CHARS, MAX_MEDIA_PROMPT_BYTES, "Audio prompt")
-    if request.headers.get("x-payg-internal") != "true":
+    if request.headers.get("x-payg-internal") != WEBSOCKET_KEY:
         await check_audio_rate_limit(request, authorization, x_client_id)
 
     url = f"https://gen.pollinations.ai/audio/{prompt}?model=acestep&key={PKEY}"
@@ -225,7 +225,7 @@ async def gentts(
 
     prompt = normalize_prompt_value(prompt, "prompt")
     enforce_prompt_size(prompt, MAX_MEDIA_PROMPT_CHARS, MAX_MEDIA_PROMPT_BYTES, "Audio prompt")
-    if request.headers.get("x-payg-internal") != "true":
+    if request.headers.get("x-payg-internal") != WEBSOCKET_KEY:
         await check_audio_rate_limit(request, authorization, x_client_id)
 
     url = f"https://gen.pollinations.ai/audio/{prompt}?model=universal-3-pro&key={PKEY3}"
@@ -316,7 +316,7 @@ async def genvideo(request: Request, prompt: str = None, authorization: str = He
     enforce_prompt_size(
         prompt, MAX_MEDIA_PROMPT_CHARS, MAX_MEDIA_PROMPT_BYTES, "Video prompt"
     )
-    if request.headers.get("x-payg-internal") != "true":
+    if request.headers.get("x-payg-internal") != WEBSOCKET_KEY:
         await check_video_rate_limit(request, authorization, x_client_id)
 
     RATIO_MAP = {
@@ -469,7 +469,7 @@ async def genvideo_airforce(
     enforce_prompt_size(
         prompt, MAX_MEDIA_PROMPT_CHARS, MAX_MEDIA_PROMPT_BYTES, "Video prompt"
     )
-    if request.headers.get("x-payg-internal") != "true":
+    if request.headers.get("x-payg-internal") != WEBSOCKET_KEY:
         await check_video_rate_limit(request, authorization, x_client_id)
 
     payload = {
@@ -634,7 +634,7 @@ async def generate_text(
             provider = "cerebras"
             chosen_model = "gpt-oss-120b"
 
-    if request.headers.get("x-payg-internal") != "true":
+    if request.headers.get("x-payg-internal") != WEBSOCKET_KEY:
         await check_chat_rate_limit(request, authorization, x_client_id)
 
     body["model"] = chosen_model