Update modules/app.py

modules/app.py

@@ -2,16 +2,25 @@ from fastapi import FastAPI, HTTPException, Request, Depends
 import aiohttp
 import base64
 import io
-import asyncio
+import os
 import random
 import string
 import json
 
 app = FastAPI()
 
+whisper_origin = os.getenv("WHISPER_ORIGIN")
+whisper_base_url = os.getenv("WHISPER_BASE_URL")
+img2location_name = os.getenv("IMG2LOCATION_NAME")
+img2location_origin = os.getenv("IMG2LOCATION_ORIGIN")
+img2location_base_url = os.getenv("IMG2LOCATION_BASE_URL")
+pixart_sigma_base_url = os.getenv("PIXART_SIGMA_BASE_URL")
+allowed_user_agent = os.getenv("ALLOWED_USER_AGENT")
+
+
 async def verify_user_agent(request: Request):
     user_agent = request.headers.get('User-Agent', None)
-    if user_agent != "Vercel Edge Functions":
+    if user_agent != allowed_user_agent:
         raise HTTPException(status_code=403, detail="Access denied")
     return True
 
@@ -40,9 +49,9 @@ async def whisper(request: Request):
         'Connection': 'keep-alive',
         'Content-Type': 'application/json',
         'DNT': '1',
-        'Origin': 'https://deepinfra.com',
+        'Origin': whisper_origin,
         'Pragma': 'no-cache',
-        'Referer': 'https://deepinfra.com/',
+        'Referer': f'{whisper_origin}/',
         'Sec-Fetch-Dest': 'empty',
         'Sec-Fetch-Mode': 'cors',
         'Sec-Fetch-Site': 'same-site',
@@ -65,7 +74,7 @@ async def whisper(request: Request):
             json_data = '{"audio": "' + audio_base64 + '"}'
 
             # Post request to the API
-            async with session.post('https://api.deepinfra.com/v1/inference/openai/whisper-large', headers=headers, data=json_data) as post_resp:
+            async with session.post(f'{whisper_base_url}/v1/inference/openai/whisper-large', headers=headers, data=json_data) as post_resp:
                 if post_resp.status != 200:
                     return f"API request failed: {post_resp.status}"
                 return await post_resp.json()
@@ -106,10 +115,10 @@ async def img2location(request: Request):
         'accept-language': 'en-US,en;q=0.9',
         'cache-control': 'no-cache',
         'dnt': '1',
-        'origin': 'https://geospy.ai',
+        'origin': img2location_origin,
         'pragma': 'no-cache',
         'priority': 'u=1, i',
-        'referer': 'https://geospy.ai/',
+        'referer': f'{img2location_origin}/',
         'sec-ch-ua': '"Chromium";v="124", "Microsoft Edge";v="124", "Not-A.Brand";v="99"',
         'sec-ch-ua-mobile': '?0',
         'sec-ch-ua-platform': '"Windows"',
@@ -131,7 +140,7 @@ async def img2location(request: Request):
         data.add_field('image', io.BytesIO(image_data), filename="image.png", content_type='image/png')
 
         # Sending the POST request
-        async with session.post('https://locate-image-7cs5mab6na-uc.a.run.app/', headers=headers, data=data) as response:
+        async with session.post(img2location_base_url, headers=headers, data=data) as response:
             if response.status != 200:
                 return f"Failed to upload image: HTTP {response.status}"
             json_response = await response.json()
@@ -145,6 +154,9 @@ async def img2location(request: Request):
                     google_maps = f"https://www.google.com/maps/search/?api=1&query={latitude},{longitude}"
                 
                     json_response['message'] += f"\n\nView on Google Maps: {google_maps}"
+                
+                if img2location_name in json_response['message'].lower():
+                    raise HTTPException(status_code=400, detail="We are not allowed to process this image. Please try another one.")
 
                 return json_response
             
@@ -202,14 +214,14 @@ async def pixart_sigma(request: Request):
     }
 
     async with aiohttp.ClientSession() as session:
-        async with session.post('https://pixart-alpha-pixart-sigma.hf.space/queue/join', params=params, headers=headers, json=json_data) as response:
+        async with session.post(f'{pixart_sigma_base_url}/queue/join', params=params, headers=headers, json=json_data) as response:
             print(response.text)
 
         params = {
             'session_hash': hash,
         }
 
-        async with session.get('https://pixart-alpha-pixart-sigma.hf.space/queue/data', params=params, headers=headers) as response:
+        async with session.get(f'{pixart_sigma_base_url}/queue/data', params=params, headers=headers) as response:
             async for line in response.content:
                 try:
                     if line: