Update main.py

main.py

@@ -17,11 +17,12 @@ load_dotenv()
 app = FastAPI()
 
 # Get API keys and secret endpoint from environment variables
-api_keys_str = os.getenv('API_KEYS')
+api_keys_str = os.getenv('API_KEYS') #deprecated -_-
 valid_api_keys = api_keys_str.split(',') if api_keys_str else []
 secret_api_endpoint = os.getenv('SECRET_API_ENDPOINT')
 secret_api_endpoint_2 = os.getenv('SECRET_API_ENDPOINT_2')
 secret_api_endpoint_3 = os.getenv('SECRET_API_ENDPOINT_3')  # New endpoint for searchgpt
+image_endpoint = os.getenv("IMAGE_ENDPOINT")
 
 # Validate if the main secret API endpoints are set
 if not secret_api_endpoint or not secret_api_endpoint_2 or not secret_api_endpoint_3:
@@ -145,7 +146,89 @@ available_model_ids = [
     "llama-3.1-8b", "gemini-1.5-flash", "mixtral-8x7b" , "command-r","gemini-pro",
     "gpt-3.5-turbo", "command"
 ]
+@app.get("/images/generations")
+async def generate_image(
+    prompt: str,
+    model: str = "flux",  # Default model
+    seed: Optional[int] = None,
+    width: Optional[int] = None,
+    height: Optional[int] = None,
+    nologo: Optional[bool] = True,
+    private: Optional[bool] = None,
+    enhance: Optional[bool] = None,
+):
+    """
+    Generate an image using the Image Generation API.
 
+    Example usage:
+    - /images/generations?prompt=beautiful+sunset&model=flux&width=1024&height=768
+    """
+    # Ensure the IMAGE_ENDPOINT is configured
+    if not image_endpoint:
+        raise HTTPException(
+            status_code=500, 
+            detail="Image endpoint not configured in environment variables."
+        )
+    
+    # Check if prompt is valid
+    if not prompt or len(prompt.strip()) == 0:
+        raise HTTPException(
+            status_code=400, 
+            detail="Prompt is required and cannot be empty. Example: prompt=beautiful+sunset"
+        )
+    
+    # Construct the URL with the prompt
+    url = f"{image_endpoint}/{prompt.strip()}"
+    
+    # Prepare query parameters
+    params = {
+        "model": model,
+        "seed": seed,
+        "width": width,
+        "height": height,
+        "nologo": nologo,
+        "private": private,
+        "enhance": enhance,
+    }
+    # Remove keys with `None` values to avoid invalid params
+    params = {k: v for k, v in params.items() if v is not None}
+    
+    try:
+        # Send GET request to the image generation endpoint
+        async with httpx.AsyncClient() as client:
+            response = await client.get(url, params=params)
+        
+        # Handle non-successful HTTP status codes
+        if response.status_code == 400:
+            raise HTTPException(
+                status_code=400, 
+                detail="Invalid request to the image generation API. Please check your prompt and parameters."
+            )
+        elif response.status_code == 500:
+            raise HTTPException(
+                status_code=500, 
+                detail="Server error occurred while generating the image. Please try again later."
+            )
+        elif response.status_code != 200:
+            raise HTTPException(
+                status_code=response.status_code, 
+                detail=f"Error generating image: {response.text}"
+            )
+        
+        # Return the image as a file response
+        return StreamingResponse(response.aiter_bytes(), media_type="image/jpeg")
+    
+    except httpx.TimeoutException:
+        raise HTTPException(
+            status_code=504, 
+            detail="The request to the image generation API timed out. Please try again later."
+        )
+    except httpx.RequestError as e:
+        # Handle other request errors
+        raise HTTPException(
+            status_code=500, 
+            detail=f"Error contacting the image endpoint: {str(e)}"
+        )
 @app.post("/chat/completions")
 @app.post("/v1/chat/completions")
 async def get_completion(payload: Payload,request: Request):