Update main.py

main.py

@@ -206,9 +206,9 @@ async def get_completion(payload: Payload, request: Request):
     
     return StreamingResponse(stream_generator(payload_dict), media_type="application/json")
 # Remove the duplicated endpoint and combine the functionality
-@app.get("/images/generations") #pollinations.ai thanks to them :)
+@app.api_route("/images/generations", methods=["GET", "POST"])  # Support both GET and POST
 async def generate_image(
-    prompt: str,
+    prompt: Optional[str] = None, 
     model: str = "flux",  # Default model
     seed: Optional[int] = None,
     width: Optional[int] = None,
@@ -216,6 +216,7 @@ async def generate_image(
     nologo: Optional[bool] = True,
     private: Optional[bool] = None,
     enhance: Optional[bool] = None,
+    request: Request = None,  # Access raw POST data
 ):
     """
     Generate an image using the Image Generation API.
@@ -223,22 +224,31 @@ async def generate_image(
     # Validate the image endpoint
     if not image_endpoint:
         raise HTTPException(status_code=500, detail="Image endpoint not configured in environment variables.")
-    
-    # Validate prompt
-    if not prompt or not prompt.strip():
-        raise HTTPException(status_code=400, detail="Prompt cannot be empty")
-        
+
+    # Handle GET and POST prompts
+    if request.method == "POST":
+        try:
+            body = await request.json()  # Parse JSON body
+            prompt = body.get("prompt", "").strip()
+            if not prompt:
+                raise HTTPException(status_code=400, detail="Prompt cannot be empty")
+        except Exception:
+            raise HTTPException(status_code=400, detail="Invalid JSON payload")
+    elif request.method == "GET":
+        if not prompt or not prompt.strip():
+            raise HTTPException(status_code=400, detail="Prompt cannot be empty")
+        prompt = prompt.strip()
+
     # Sanitize and encode the prompt
-    sanitized_prompt = prompt.strip()
-    encoded_prompt = httpx.QueryParams({'prompt': sanitized_prompt}).get('prompt')
-    
+    encoded_prompt = httpx.QueryParams({'prompt': prompt}).get('prompt')
+
     # 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):
@@ -253,12 +263,12 @@ async def generate_image(
         params['private'] = str(private).lower()
     if enhance is not None:
         params['enhance'] = str(enhance).lower()
-    
+
     try:
         timeout = httpx.Timeout(60.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")
@@ -271,7 +281,7 @@ async def generate_image(
                     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/'):
@@ -279,7 +289,7 @@ async def generate_image(
                     status_code=500,
                     detail=f"Unexpected content type received: {content_type}"
                 )
-            
+
             return StreamingResponse(
                 response.iter_bytes(),
                 media_type=content_type,
@@ -288,7 +298,7 @@ async def generate_image(
                     'Pragma': 'no-cache'
                 }
             )
-            
+
     except httpx.TimeoutException:
         raise HTTPException(status_code=504, detail="Image generation request timed out")
     except httpx.RequestError as e: