Update main.py

main.py

@@ -15,11 +15,21 @@ from io import BytesIO
 from PIL import Image
 import requests
 import os
+import logging
 from dotenv import load_dotenv
 
 # Load environment variables
 load_dotenv()
 
+# Validate environment variables
+assert os.getenv('MONGO_USER') and os.getenv('MONGO_PASSWORD') and os.getenv('MONGO_HOST'), "MongoDB credentials missing!"
+assert os.getenv('LLM_API_KEY'), "LLM API Key missing!"
+assert os.getenv('BFL_API_KEY'), "BFL API Key missing!"
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 # Set the Hugging Face cache directory to a writable location
 os.environ['HF_HOME'] = '/tmp/huggingface_cache'
 
@@ -51,16 +61,16 @@ aura_sr = None
 @app.on_event("startup")
 async def startup():
     global llm, aura_sr
-    llm = ChatGroq(
-        model="llama-3.3-70b-versatile",
-        temperature=0.7,
-        max_tokens=1024,
-        api_key=os.getenv('LLM_API_KEY'),
-    )
     try:
+        llm = ChatGroq(
+            model="llama-3.3-70b-versatile",
+            temperature=0.7,
+            max_tokens=1024,
+            api_key=os.getenv('LLM_API_KEY'),
+        )
         aura_sr = AuraSR.from_pretrained("fal/AuraSR-v2")
     except Exception as e:
-        print(f"Error initializing AuraSR: {e}")
+        logger.error(f"Error initializing models: {e}")
 
 @app.on_event("shutdown")
 def shutdown():
@@ -104,11 +114,13 @@ def save_image_locally(image, filename):
     return filepath
 
 def fetch_image(url):
-    with requests.Session() as session:
-        response = session.get(url, timeout=10)
-        if response.status_code != 200:
-            raise HTTPException(status_code=400, detail="Failed to fetch image.")
-        return Image.open(BytesIO(response.content))
+    try:
+        with requests.Session() as session:
+            response = session.get(url, timeout=10)
+            response.raise_for_status()
+            return Image.open(BytesIO(response.content))
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Error fetching image: {str(e)}")
 
 def poll_for_image_result(request_id, headers):
     timeout = 60
@@ -131,7 +143,6 @@ def poll_for_image_result(request_id, headers):
 
     raise HTTPException(status_code=500, detail="Image generation timed out.")
 
-# Endpoints
 @app.post("/new-chat", response_model=dict)
 async def new_chat():
     chat_id = generate_chat_id()
@@ -176,14 +187,12 @@ async def generate_image(request: ImageRequest):
         "max_sequence_length": 512,
     }
 
-    with requests.Session() as session:
-        response = session.post("https://api.bfl.ml/v1/flux-pro-1.1", headers=headers, json=payload, timeout=10).json()
+    response = make_request_with_retries("https://api.bfl.ml/v1/flux-pro-1.1", headers, payload)
 
     if "id" not in response:
         raise HTTPException(status_code=500, detail="Error generating image: ID missing from response")
 
     image_url = poll_for_image_result(response["id"], headers)
-
     image = fetch_image(image_url)
     filename = f"generated_{uuid.uuid4()}.png"
     filepath = save_image_locally(image, filename)
@@ -199,20 +208,21 @@ async def upscale_image(request: UpscaleRequest):
     if aura_sr is None:
         raise HTTPException(status_code=500, detail="Upscaling model not initialized.")
 
-    try:
-        img = fetch_image(request.image_url)
-        upscaled_image = aura_sr.upscale_4x_overlapped(img)
+    img = await run_in_threadpool(fetch_image, request.image_url)
 
+    def perform_upscaling():
+        upscaled_image = aura_sr.upscale_4x_overlapped(img)
         filename = f"upscaled_{uuid.uuid4()}.png"
-        filepath = save_image_locally(upscaled_image, filename)
+        return save_image_locally(upscaled_image, filename)
 
-        return {
-            "status": "Upscaling successful",
-            "file_path": filepath,
-            "file_url": f"/images/{filename}",
-        }
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error during upscaling: {str(e)}")
+    future = executor.submit(perform_upscaling)
+    filepath = await run_in_threadpool(lambda: future.result())
+
+    return {
+        "status": "Upscaling successful",
+        "file_path": filepath,
+        "file_url": f"/images/{os.path.basename(filepath)}",
+    }
 
 @app.get("/")
 async def root():