Update main.py

main.py

@@ -1,26 +1,36 @@
 import os
 import re
-from dotenv import load_dotenv
-from fastapi import FastAPI, HTTPException, Request, Depends, Security, Query
-from fastapi.responses import StreamingResponse, HTMLResponse, JSONResponse, FileResponse, PlainTextResponse
-from fastapi.security import APIKeyHeader
-from pydantic import BaseModel
-import httpx
-from functools import lru_cache
-from pathlib import Path
 import json
-import datetime
 import time
+import datetime
 import asyncio
-from starlette.status import HTTP_403_FORBIDDEN
+import hashlib
+from pathlib import Path
+from functools import lru_cache
+from typing import List, Dict, Any, Optional, Union
+
+import httpx
+import requests
 import cloudscraper
 from concurrent.futures import ThreadPoolExecutor
-import uvloop
+
+from dotenv import load_dotenv
+from fastapi import FastAPI, HTTPException, Request, Depends, Security, Query, APIRouter
+from fastapi.responses import (
+    StreamingResponse,
+    HTMLResponse,
+    JSONResponse,
+    FileResponse,
+    PlainTextResponse,
+)
+from fastapi.security import APIKeyHeader
 from fastapi.middleware.gzip import GZipMiddleware
 from starlette.middleware.cors import CORSMiddleware
+from starlette.status import HTTP_403_FORBIDDEN
+
+import uvloop
 import contextlib
-import requests
-from typing import List, Dict, Any, Optional, Union # Import Optional and other typing helpers
+
 
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
@@ -43,6 +53,9 @@ app.add_middleware(
     allow_methods=["*"],
     allow_headers=["*"],
 )
+@lru_cache(maxsize=256)
+def cached_url(url: str):
+    return url
 
 @lru_cache(maxsize=1)
 def get_env_vars():
@@ -660,52 +673,45 @@ async def get_completion(payload: Payload, request: Request, authenticated: bool
         except Exception as e:
             print(f"An unexpected error occurred during non-streaming chat completion proxy: {e}")
             raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {str(e)}")
-
-@app.post("/images/generations")
-async def create_image(payload: ImageGenerationPayload, request: Request, authenticated: bool = Depends(verify_api_key)):
+router = APIRouter()
+
+@router.get("/images/{prompt:path}")
+async def create_image(
+    prompt: str,
+    request: Request,
+    authenticated: bool = Depends(verify_api_key)
+):
     if not server_status:
-        raise HTTPException(
-            status_code=503,
-            content={"message": "Server is under maintenance. Please try again later."}
-        )
+        raise HTTPException(status_code=503, detail="Server is under maintenance.")
 
-    if payload.model not in supported_image_models:
-        raise HTTPException(
-            status_code=400,
-            detail=f"Model '{payload.model}' is not supported for image generation. Supported models are: {', '.join(supported_image_models)}"
-        )
+    # forward all GET params
+    query = request.url.query
 
-    usage_tracker.record_request(request=request, model=payload.model, endpoint="/images/generations")
+    base = "https://image.pollinations.ai/prompt/"
+    final_url = f"{base}{prompt}"
+    if query:
+        final_url += f"?{query}"
 
-    api_payload = {
-        "model": payload.model,
-        "prompt": payload.prompt,
-        "size": payload.size,
-        "n": payload.number
-    }
-
-    target_api_url = get_env_vars().get('new_img')
-    if not target_api_url:
-        raise HTTPException(status_code=500, detail="Image generation API endpoint (NEW_IMG) not configured.")
+    # caching
+    final_url = cached_url(final_url)
 
     try:
-        client = get_async_client()
-        response = await client.post(target_api_url, json=api_payload)
+        client = httpx.AsyncClient(timeout=60)
+        resp = await client.get(final_url)
 
-        response.raise_for_status()
+        if resp.status_code != 200:
+            raise HTTPException(status_code=resp.status_code, detail="Image generation failed.")
 
-        return JSONResponse(content=response.json())
+        return StreamingResponse(
+            resp.aiter_bytes(),
+            media_type="image/jpeg"
+        )
 
     except httpx.TimeoutException:
-        raise HTTPException(status_code=504, detail="Image generation request timed out.")
-    except httpx.RequestError as e:
-        raise HTTPException(status_code=502, detail=f"Error connecting to image generation service: {e}")
-    except httpx.HTTPStatusError as e:
-        error_detail = e.response.json().get("detail", f"Image generation failed with status code: {e.response.status_code}")
-        raise HTTPException(status_code=e.response.status_code, detail=error_detail)
+        raise HTTPException(status_code=504, detail="Image generation timeout.")
+
     except Exception as e:
-        print(f"An unexpected error occurred during image generation: {e}")
-        raise HTTPException(status_code=500, detail=f"An unexpected error occurred during image generation: {e}")
+        raise HTTPException(status_code=500, detail=f"Unexpected error: {e}")
 
 @app.get("/usage")
 async def get_usage_json(days: int = 7):