Update main.py

main.py

@@ -5,7 +5,7 @@ from fastapi.concurrency import run_in_threadpool
 from fastapi.middleware.cors import CORSMiddleware
 import uuid
 import time
-import tempfile
+tempfile
 from concurrent.futures import ThreadPoolExecutor
 from pymongo import MongoClient
 from urllib.parse import quote_plus
@@ -15,7 +15,10 @@ from io import BytesIO
 from PIL import Image
 import requests
 import os
-import os
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
 
 # Set the Hugging Face cache directory to a writable location
 os.environ['HF_HOME'] = '/tmp/huggingface_cache'
@@ -31,22 +34,19 @@ app.add_middleware(
 )
 
 # Globals
-executor = ThreadPoolExecutor(max_workers=10)
+executor = ThreadPoolExecutor(max_workers=5)
 llm = None
-client = MongoClient(f"mongodb+srv://hammad:{quote_plus('momimaad@123')}@cluster0.2a9yu.mongodb.net/")
-db = client["Flux"]
+mongo_client = MongoClient(f"mongodb+srv://{os.getenv('MONGO_USER')}:{quote_plus(os.getenv('MONGO_PASSWORD'))}@{os.getenv('MONGO_HOST')}/")
+db = mongo_client["Flux"]
 collection = db["chat_histories"]
 chat_sessions = {}
 
-# Use a temporary directory for storing images
+# Temporary directory for storing images
 image_storage_dir = tempfile.mkdtemp()
-print(f"Temporary directory for images: {image_storage_dir}")
-
-# Serve the temporary image directory as a static file directory
 app.mount("/images", StaticFiles(directory=image_storage_dir), name="images")
 
-# Dictionary to store images temporarily
-image_store = {}
+# Initialize AuraSR during startup
+aura_sr = None
 
 @app.on_event("startup")
 async def startup():
@@ -55,19 +55,16 @@ async def startup():
         model="llama-3.3-70b-versatile",
         temperature=0.7,
         max_tokens=1024,
-        api_key="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",
+        api_key=os.getenv('LLM_API_KEY'),
     )
-    # Ensure AuraSR is initialized properly
     try:
         aura_sr = AuraSR.from_pretrained("fal/AuraSR-v2")
     except Exception as e:
         print(f"Error initializing AuraSR: {e}")
-        aura_sr = None
-
 
 @app.on_event("shutdown")
 def shutdown():
-    client.close()
+    mongo_client.close()
     executor.shutdown()
 
 # Pydantic models
@@ -85,7 +82,10 @@ class ImageRequest(BaseModel):
     user_prompt: str
     chat_id: str
 
-# Helper Functions
+class UpscaleRequest(BaseModel):
+    image_url: str
+
+# Helper functions
 def generate_chat_id():
     chat_id = str(uuid.uuid4())
     chat_sessions[chat_id] = collection
@@ -103,6 +103,34 @@ def save_image_locally(image, filename):
     image.save(filepath, format="PNG")
     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))
+
+def poll_for_image_result(request_id, headers):
+    timeout = 60
+    start_time = time.time()
+
+    while time.time() - start_time < timeout:
+        time.sleep(0.5)
+        with requests.Session() as session:
+            result = session.get(
+                "https://api.bfl.ml/v1/get_result",
+                headers=headers,
+                params={"id": request_id},
+                timeout=10
+            ).json()
+
+        if result["status"] == "Ready":
+            return result["result"].get("sample")
+        elif result["status"] == "Error":
+            raise HTTPException(status_code=500, detail=f"Image generation failed: {result.get('error', 'Unknown error')}")
+
+    raise HTTPException(status_code=500, detail="Image generation timed out.")
+
 # Endpoints
 @app.post("/new-chat", response_model=dict)
 async def new_chat():
@@ -111,148 +139,81 @@ async def new_chat():
 
 @app.post("/generate-image", response_model=dict)
 async def generate_image(request: ImageRequest):
-    def process_request():
-        chat_history = get_chat_history(request.chat_id)
-        prompt = f"""
-        You are a professional assistant responsible for crafting a clear and visually compelling prompt for an image generation model. Your task is to generate a high-quality prompt for creating both the **main subject** and the **background** of the image.
-
-        Image Specifications:
-        - **Subject**: Focus on **{request.subject}**, highlighting its defining features, expressions, and textures.
-        - **Style**: Emphasize the **{request.style}**, capturing its key characteristics.
-        - **Background**: Create a background with **{request.background_details}** that complements and enhances the subject. Ensure it aligns with the color theme and overall composition.
-        - **Camera and Lighting**:
-          - Lighting: Apply **{request.lighting_conditions}**, emphasizing depth, highlights, and shadows to accentuate the subject and harmonize the background.
-        - **Framing**: Use a **{request.framing_style}** to enhance the composition around both the subject and the background.
-        - **Materials**: Highlight textures like **{request.material_details}**, with realistic details and natural imperfections on the subject and background.
-        - **Key Elements**: Include **{request.elements}** to enrich the subject’s details and add cohesive elements to the background.
-        - **Color Theme**: Follow the **{request.color_theme}** to set the mood and tone for the entire scene.
-        - Negative Prompt: Avoid grainy, blurry, or deformed outputs.
-        - **Text to Include in Image**: Clearly display the text **"{request.text}"** as part of the composition (e.g., on a card, badge, or banner) attached to the subject in a realistic and contextually appropriate way.
-        """
-        refined_prompt = llm.invoke(prompt).content.strip()
-
-        # Save the prompt in MongoDB
-        collection.insert_one({"session_id": request.chat_id, "role": "user", "content": request.user_prompt})
-        collection.insert_one({"session_id": request.chat_id, "role": "ai", "content": refined_prompt})
-
-        # API call to image generation service
-        url = "https://api.bfl.ml/v1/flux-pro-1.1"
-        headers = {
-            "accept": "application/json",
-            "x-key": "4f69d408-4979-4812-8ad2-ec6d232c9ddf",
-            "Content-Type": "application/json"
-        }
-        payload = {
-            "prompt": refined_prompt,
-            "width": 1024,
-            "height": 1024,
-            "guidance_scale": 1,
-            "num_inference_steps": 50,
-            "max_sequence_length": 512,
-        }
+    chat_history = get_chat_history(request.chat_id)
+    prompt = f"""
+    You are a professional assistant responsible for crafting a clear and visually compelling prompt for an image generation model. Your task is to generate a high-quality prompt for creating both the **main subject** and the **background** of the image.
+
+    Image Specifications:
+    - **Subject**: Focus on **{request.subject}**, highlighting its defining features, expressions, and textures.
+    - **Style**: Emphasize the **{request.style}**, capturing its key characteristics.
+    - **Background**: Create a background with **{request.background_details}** that complements and enhances the subject. Ensure it aligns with the color theme and overall composition.
+    - **Camera and Lighting**:
+      - Lighting: Apply **{request.lighting_conditions}**, emphasizing depth, highlights, and shadows to accentuate the subject and harmonize the background.
+    - **Framing**: Use a **{request.framing_style}** to enhance the composition around both the subject and the background.
+    - **Materials**: Highlight textures like **{request.material_details}**, with realistic details and natural imperfections on the subject and background.
+    - **Key Elements**: Include **{request.elements}** to enrich the subject’s details and add cohesive elements to the background.
+    - **Color Theme**: Follow the **{request.color_theme}** to set the mood and tone for the entire scene.
+    - Negative Prompt: Avoid grainy, blurry, or deformed outputs.
+    - **Text to Include in Image**: Clearly display the text **"{request.text}"** as part of the composition (e.g., on a card, badge, or banner) attached to the subject in a realistic and contextually appropriate way.
+    """
+
+    refined_prompt = llm.invoke(prompt).content.strip()
+    collection.insert_one({"session_id": request.chat_id, "role": "user", "content": request.user_prompt})
+    collection.insert_one({"session_id": request.chat_id, "role": "ai", "content": refined_prompt})
+
+    headers = {
+        "accept": "application/json",
+        "x-key": os.getenv('BFL_API_KEY'),
+        "Content-Type": "application/json"
+    }
 
-        response = requests.post(url, headers=headers, json=payload).json()
-        if "id" not in response:
-            raise HTTPException(status_code=500, detail="Error generating image: ID missing from response")
+    payload = {
+        "prompt": refined_prompt,
+        "width": 1024,
+        "height": 1024,
+        "guidance_scale": 1,
+        "num_inference_steps": 50,
+        "max_sequence_length": 512,
+    }
 
-        request_id = response["id"]
+    with requests.Session() as session:
+        response = session.post("https://api.bfl.ml/v1/flux-pro-1.1", headers=headers, json=payload, timeout=10).json()
 
-        # Poll for the image result
-        while True:
-            time.sleep(0.5)
-            result = requests.get(
-                "https://api.bfl.ml/v1/get_result",
-                headers=headers,
-                params={"id": request_id},
-            ).json()
+    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)
 
-            if result["status"] == "Ready":
-                if "result" in result and "sample" in result["result"]:
-                    image_url = result["result"]["sample"]
-
-                    # Download the image
-                    image_response = requests.get(image_url)
-                    if image_response.status_code == 200:
-                        img = Image.open(BytesIO(image_response.content))
-                        filename = f"generated_{uuid.uuid4()}.png"
-                        filepath = save_image_locally(img, filename)
-
-                        # Generate a URL for the image
-                        file_url = f"/images/{filename}"
-                        return filepath, file_url
-                    else:
-                        raise HTTPException(status_code=500, detail="Failed to download the image")
-                else:
-                    raise HTTPException(status_code=500, detail="Expected 'sample' key not found in the result")
-            elif result["status"] == "Error":
-                raise HTTPException(status_code=500, detail=f"Image generation failed: {result.get('error', 'Unknown error')}")
-
-    future = executor.submit(process_request)
-    filepath, file_url = await run_in_threadpool(future.result)
+    image = fetch_image(image_url)
+    filename = f"generated_{uuid.uuid4()}.png"
+    filepath = save_image_locally(image, filename)
 
     return {
         "status": "Image generated successfully",
         "file_path": filepath,
-        "file_url": file_url,
+        "file_url": f"/images/{filename}",
     }
 
-class UpscaleRequest(BaseModel):
-    image_url: str
-
-
-def process_image(image_url):
-    if aura_sr is None:
-        raise RuntimeError("aura_sr is None. Ensure it's initialized during startup.")
-
-    response = requests.get(image_url)
-    img = Image.open(BytesIO(response.content))
-
-    try:
-        upscaled_image = aura_sr.upscale_4x_overlapped(img)
-    except Exception as e:
-        raise RuntimeError(f"Error during upscaling: {str(e)}")
-    
-    filename = f"upscaled_{uuid.uuid4()}.png"
-    filepath = save_image_locally(upscaled_image, filename)
-    return filepath
-
-
-
 @app.post("/upscale-image", response_model=dict)
 async def upscale_image(request: UpscaleRequest):
     if aura_sr is None:
         raise HTTPException(status_code=500, detail="Upscaling model not initialized.")
 
     try:
-        # Fetch the image from the provided URL
-        response = requests.get(request.image_url)
-        if response.status_code != 200:
-            raise HTTPException(status_code=400, detail="Failed to fetch the image from the provided URL.")
-        
-        img = Image.open(BytesIO(response.content))
-
-        # Perform upscaling
+        img = fetch_image(request.image_url)
         upscaled_image = aura_sr.upscale_4x_overlapped(img)
 
-        # Save the upscaled image locally
         filename = f"upscaled_{uuid.uuid4()}.png"
         filepath = save_image_locally(upscaled_image, filename)
 
-        # Generate a public URL for the upscaled image
-        file_url = f"/images/{filename}"
-
         return {
             "status": "Upscaling successful",
             "file_path": filepath,
-            "file_url": file_url,
+            "file_url": f"/images/{filename}",
         }
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error during upscaling: {str(e)}")
 
-
-
-
-
 @app.get("/")
 async def root():
-    return {"message": "API is up and running!"}
+    return {"message": "API is up and running!"}