Update main.py

main.py

@@ -146,89 +146,6 @@ 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):
@@ -284,6 +201,111 @@ async def get_completion(payload: Payload,request: Request):
             raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {e}")
 
     return StreamingResponse(stream_generator(payload_dict), media_type="application/json")
+@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] = None,
+    private: Optional[bool] = None,
+    enhance: Optional[bool] = None,
+):
+    """
+    Generate an image using the Image Generation API.
+    """
+    # Ensure the IMAGE_ENDPOINT is configured
+    image_endpoint = os.getenv("IMAGE_ENDPOINT", "https://image.pollinations.ai/prompt")
+    if not image_endpoint:
+        raise HTTPException(status_code=500, detail="Image endpoint not configured in environment variables.")
+    
+    # Construct the URL with the prompt
+    url = f"{image_endpoint}/{prompt}"
+    
+    # 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)
+        
+        # If the response is not successful, raise an error
+        if 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.iter_content(), media_type="image/jpeg")
+    
+    except httpx.RequestError as e:
+        # Handle request errors
+        raise HTTPException(status_code=500, detail=f"Error contacting the image endpoint: {str(e)}")
+@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.
+    """
+    # Ensure the IMAGE_ENDPOINT is configured
+    if not image_endpoint:
+        raise HTTPException(status_code=500, detail="Image endpoint not configured in environment variables.")
+    
+    # Construct the URL with the prompt
+    url = f"{image_endpoint}/{prompt}"
+    
+    # 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)
+        
+        # If the response is not successful, raise an error
+        if 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.iter_content(), media_type="image/jpeg")
+    
+    except httpx.RequestError as e:
+        # Handle request errors
+        raise HTTPException(status_code=500, detail=f"Error contacting the image endpoint: {str(e)}")
 @app.get("/playground", response_class=HTMLResponse)
 async def playground():
     # Open and read the content of playground.html (in the same folder as the app)