Update main.py

main.py

@@ -201,7 +201,8 @@ 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") #pollinations.ai :) thanks to them check their website
+# Remove the duplicated endpoint and combine the functionality
+@app.get("/images/generations") #pollinations.ai thanks to them :)
 async def generate_image(
     prompt: str,
     model: str = "flux",  # Default model
@@ -215,44 +216,81 @@ async def generate_image(
     """
     Generate an image using the Image Generation API.
     """
-    # Ensure the IMAGE_ENDPOINT is configured
+    # Validate the image endpoint
     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}"
+    # Validate prompt
+    if not prompt or not prompt.strip():
+        raise HTTPException(status_code=400, detail="Prompt cannot be empty")
+        
+    # Sanitize and encode the prompt
+    sanitized_prompt = prompt.strip()
+    encoded_prompt = httpx.QueryParams({'prompt': sanitized_prompt}).get('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}
+    # Construct the URL with the encoded prompt
+    base_url = image_endpoint.rstrip('/')  # Remove trailing slash if present
+    url = f"{base_url}/{encoded_prompt}"
+    
+    # Prepare query parameters with validation
+    params = {}
+    
+    if model and isinstance(model, str):
+        params['model'] = model
+    if seed is not None and isinstance(seed, int):
+        params['seed'] = seed
+    if width is not None and isinstance(width, int) and 64 <= width <= 2048:
+        params['width'] = width
+    if height is not None and isinstance(height, int) and 64 <= height <= 2048:
+        params['height'] = height
+    if nologo is not None:
+        params['nologo'] = str(nologo).lower()
+    if private is not None:
+        params['private'] = str(private).lower()
+    if enhance is not None:
+        params['enhance'] = str(enhance).lower()
     
     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}"
+        timeout = httpx.Timeout(30.0)  # Set a reasonable timeout
+        async with httpx.AsyncClient(timeout=timeout) as client:
+            response = await client.get(url, params=params, follow_redirects=True)
+            
+            # Check for various error conditions
+            if response.status_code == 404:
+                raise HTTPException(status_code=404, detail="Image generation service not found")
+            elif response.status_code == 400:
+                raise HTTPException(status_code=400, detail="Invalid parameters provided to image service")
+            elif response.status_code == 429:
+                raise HTTPException(status_code=429, detail="Too many requests to image service")
+            elif response.status_code != 200:
+                raise HTTPException(
+                    status_code=response.status_code,
+                    detail=f"Image generation failed with status code {response.status_code}"
+                )
+            
+            # Verify content type
+            content_type = response.headers.get('content-type', '')
+            if not content_type.startswith('image/'):
+                raise HTTPException(
+                    status_code=500,
+                    detail=f"Unexpected content type received: {content_type}"
+                )
+            
+            return StreamingResponse(
+                response.iter_bytes(),
+                media_type=content_type,
+                headers={
+                    'Cache-Control': 'no-cache',
+                    'Pragma': 'no-cache'
+                }
             )
-        
-        # Return the image as a file response
-        return StreamingResponse(response.iter_content(), media_type="image/jpeg")
-    
+            
+    except httpx.TimeoutException:
+        raise HTTPException(status_code=504, detail="Image generation request timed out")
     except httpx.RequestError as e:
-        # Handle request errors
-        raise HTTPException(status_code=500, detail=f"Error contacting the image endpoint: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Failed to contact image service: {str(e)}")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unexpected error during image generation: {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)